Skip to content

Commit 8a252aa

Browse files
xrmxanuraaga
xrmx
and
anuraaga
authored
opentelemetry-sdk: add experimental composable rule based sampler (open-telemetry#4882)
* opentelemetry-sdk: add experimental composable rule based sampler * Add CHANGELOG * Use a protocol for PredicateT * Update opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/_rule_based.py Co-authored-by: Anuraag (Rag) Agrawal <anuraaga@gmail.com> * Extend Predicate protocol to require __str__ and use it on sampler get_description While at it also provide AttributePredicate * Use a generator for getting predicates and descriptions * Add tests that matches with more than one rule --------- Co-authored-by: Anuraag (Rag) Agrawal <anuraaga@gmail.com>
1 parent f92c38a commit 8a252aa

4 files changed

Lines changed: 310 additions & 0 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
3030
([#4806](https://github.com/open-telemetry/opentelemetry-python/pull/4806))
3131
- Prevent possible endless recursion from happening in `SimpleLogRecordProcessor.on_emit`,
3232
([#4799](https://github.com/open-telemetry/opentelemetry-python/pull/4799)) and ([#4867](https://github.com/open-telemetry/opentelemetry-python/pull/4867)).
33+
- Add experimental composable rule based sampler
34+
([#4882](https://github.com/open-telemetry/opentelemetry-python/pull/4882))
3335
- Make ConcurrentMultiSpanProcessor fork safe
3436
([#4862](https://github.com/open-telemetry/opentelemetry-python/pull/4862))
3537
- `opentelemetry-exporter-otlp-proto-http`: fix retry logic and error handling for connection failures in trace, metric, and log exporters

‎opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
"composable_always_off",
1919
"composable_always_on",
2020
"composable_parent_threshold",
21+
"composable_rule_based",
2122
"composable_traceid_ratio_based",
2223
"composite_sampler",
2324
]
@@ -27,5 +28,6 @@
2728
from ._always_on import composable_always_on
2829
from ._composable import ComposableSampler, SamplingIntent
2930
from ._parent_threshold import composable_parent_threshold
31+
from ._rule_based import composable_rule_based
3032
from ._sampler import composite_sampler
3133
from ._traceid_ratio import composable_traceid_ratio_based

‎opentelemetry-sdk/src/opentelemetry/sdk/trace/_sampling_experimental/_rule_based.py‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from typing import Protocol, Sequence
18+
19+
from opentelemetry.context import Context
20+
from opentelemetry.trace import Link, SpanKind, TraceState
21+
from opentelemetry.util.types import AnyValue, Attributes
22+
23+
from ._composable import ComposableSampler, SamplingIntent
24+
from ._util import INVALID_THRESHOLD
25+
26+
27+
class PredicateT(Protocol):
28+
def __call__(
29+
self,
30+
parent_ctx: Context | None,
31+
name: str,
32+
span_kind: SpanKind | None,
33+
attributes: Attributes,
34+
links: Sequence[Link] | None,
35+
trace_state: TraceState | None,
36+
) -> bool: ...
37+
38+
def __str__(self) -> str: ...
39+
40+
41+
class AttributePredicate:
42+
"""An exact match of an attribute value"""
43+
44+
def __init__(self, key: str, value: AnyValue):
45+
self.key = key
46+
self.value = value
47+
48+
def __call__(
49+
self,
50+
parent_ctx: Context | None,
51+
name: str,
52+
span_kind: SpanKind | None,
53+
attributes: Attributes,
54+
links: Sequence[Link] | None,
55+
trace_state: TraceState | None,
56+
) -> bool:
57+
if not attributes:
58+
return False
59+
return attributes.get(self.key) == self.value
60+
61+
def __str__(self):
62+
return f"{self.key}={self.value}"
63+
64+
65+
RulesT = Sequence[tuple[PredicateT, ComposableSampler]]
66+
67+
_non_sampling_intent = SamplingIntent(
68+
threshold=INVALID_THRESHOLD, threshold_reliable=False
69+
)
70+
71+
72+
class _ComposableRuleBased(ComposableSampler):
73+
def __init__(self, rules: RulesT):
74+
# work on an internal copy of the rules
75+
self._rules = list(rules)
76+
77+
def sampling_intent(
78+
self,
79+
parent_ctx: Context | None,
80+
name: str,
81+
span_kind: SpanKind | None,
82+
attributes: Attributes,
83+
links: Sequence[Link] | None,
84+
trace_state: TraceState | None = None,
85+
) -> SamplingIntent:
86+
for predicate, sampler in self._rules:
87+
if predicate(
88+
parent_ctx=parent_ctx,
89+
name=name,
90+
span_kind=span_kind,
91+
attributes=attributes,
92+
links=links,
93+
trace_state=trace_state,
94+
):
95+
return sampler.sampling_intent(
96+
parent_ctx=parent_ctx,
97+
name=name,
98+
span_kind=span_kind,
99+
attributes=attributes,
100+
links=links,
101+
trace_state=trace_state,
102+
)
103+
return _non_sampling_intent
104+
105+
def get_description(self) -> str:
106+
rules_str = ",".join(
107+
f"({predicate}:{sampler.get_description()})"
108+
for predicate, sampler in self._rules
109+
)
110+
return f"ComposableRuleBased{{[{rules_str}]}}"
111+
112+
113+
def composable_rule_based(
114+
rules: RulesT,
115+
) -> ComposableSampler:
116+
"""Returns a consistent sampler that:
117+
118+
- Evaluates a series of rules based on predicates and returns the SamplingIntent from the first matching sampler
119+
- If no rules match, returns a non-sampling intent
120+
121+
Args:
122+
rules: A list of (Predicate, ComposableSampler) pairs, where Predicate is a function that evaluates whether a rule applies
123+
"""
124+
return _ComposableRuleBased(rules)

‎opentelemetry-sdk/tests/trace/composite_sampler/test_rule_based.py‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from opentelemetry.sdk.trace._sampling_experimental import (
16+
composable_always_off,
17+
composable_always_on,
18+
composable_rule_based,
19+
composite_sampler,
20+
)
21+
from opentelemetry.sdk.trace._sampling_experimental._rule_based import (
22+
AttributePredicate,
23+
)
24+
from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
25+
from opentelemetry.sdk.trace.sampling import Decision
26+
27+
28+
class NameIsFooPredicate:
29+
def __call__(
30+
self,
31+
parent_ctx,
32+
name,
33+
span_kind,
34+
attributes,
35+
links,
36+
trace_state,
37+
):
38+
return name == "foo"
39+
40+
def __str__(self):
41+
return "NameIsFooPredicate"
42+
43+
44+
def test_description_with_no_rules():
45+
assert (
46+
composable_rule_based(rules=[]).get_description()
47+
== "ComposableRuleBased{[]}"
48+
)
49+
50+
51+
def test_description_with_rules():
52+
rules = [
53+
(AttributePredicate("foo", "bar"), composable_always_on()),
54+
(NameIsFooPredicate(), composable_always_off()),
55+
]
56+
assert (
57+
composable_rule_based(rules=rules).get_description()
58+
== "ComposableRuleBased{[(foo=bar:ComposableAlwaysOn),(NameIsFooPredicate:ComposableAlwaysOff)]}"
59+
)
60+
61+
62+
def test_sampling_intent_match():
63+
rules = [
64+
(NameIsFooPredicate(), composable_always_on()),
65+
]
66+
assert (
67+
composable_rule_based(rules=rules)
68+
.sampling_intent(None, "foo", None, {}, None, None)
69+
.threshold
70+
== 0
71+
)
72+
73+
74+
def test_sampling_intent_no_match():
75+
rules = [
76+
(NameIsFooPredicate(), composable_always_on()),
77+
]
78+
assert (
79+
composable_rule_based(rules=rules)
80+
.sampling_intent(None, "test", None, {}, None, None)
81+
.threshold
82+
== -1
83+
)
84+
85+
86+
def test_should_sample_match():
87+
rules = [
88+
(NameIsFooPredicate(), composable_always_on()),
89+
]
90+
sampler = composite_sampler(composable_rule_based(rules=rules))
91+
92+
res = sampler.should_sample(
93+
None,
94+
RandomIdGenerator().generate_trace_id(),
95+
"foo",
96+
None,
97+
None,
98+
None,
99+
None,
100+
)
101+
102+
assert res.decision == Decision.RECORD_AND_SAMPLE
103+
assert res.trace_state is not None
104+
assert res.trace_state.get("ot", "") == "th:0"
105+
106+
107+
def test_should_sample_match_multiple_rules():
108+
rules = [
109+
(AttributePredicate("foo", "bar"), composable_always_off()),
110+
(NameIsFooPredicate(), composable_always_on()),
111+
]
112+
sampler = composite_sampler(composable_rule_based(rules=rules))
113+
114+
res = sampler.should_sample(
115+
None,
116+
RandomIdGenerator().generate_trace_id(),
117+
"foo",
118+
None,
119+
None,
120+
None,
121+
None,
122+
)
123+
124+
assert res.decision == Decision.RECORD_AND_SAMPLE
125+
assert res.trace_state is not None
126+
assert res.trace_state.get("ot", "") == "th:0"
127+
128+
129+
def test_should_sample_no_match():
130+
rules = [
131+
(NameIsFooPredicate(), composable_always_on()),
132+
]
133+
sampler = composite_sampler(composable_rule_based(rules=rules))
134+
135+
res = sampler.should_sample(
136+
None,
137+
RandomIdGenerator().generate_trace_id(),
138+
"test",
139+
None,
140+
None,
141+
None,
142+
None,
143+
)
144+
145+
assert res.decision == Decision.DROP
146+
assert res.trace_state is None
147+
148+
149+
def test_attribute_predicate_no_attributes():
150+
rules = [
151+
(AttributePredicate("foo", "bar"), composable_always_on()),
152+
]
153+
assert (
154+
composable_rule_based(rules=rules)
155+
.sampling_intent(None, "span", None, None, None, None)
156+
.threshold
157+
== -1
158+
)
159+
160+
161+
def test_attribute_predicate_no_match():
162+
rules = [
163+
(AttributePredicate("foo", "bar"), composable_always_on()),
164+
]
165+
assert (
166+
composable_rule_based(rules=rules)
167+
.sampling_intent(None, "span", None, {"foo": "foo"}, None, None)
168+
.threshold
169+
== -1
170+
)
171+
172+
173+
def test_attribute_predicate_match():
174+
rules = [
175+
(AttributePredicate("foo", "bar"), composable_always_on()),
176+
]
177+
assert (
178+
composable_rule_based(rules=rules)
179+
.sampling_intent(None, "span", None, {"foo": "bar"}, None, None)
180+
.threshold
181+
== 0
182+
)

0 commit comments

Comments
 (0)