Skip to content

Commit d5e51a4

Browse files
lzchenalrex
lzchen
and
alrex
authored
Pass in context instead of parentcontext into samplers (open-telemetry#1267)
Co-authored-by: alrex <aboten@lightstep.com>
1 parent 27dc691 commit d5e51a4

File tree

4 files changed

+96
-54
lines changed

4 files changed

+96
-54
lines changed

‎opentelemetry-sdk/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
([#1265](https://github.com/open-telemetry/opentelemetry-python/pull/1265))
99
- Allow None in sequence attributes values
1010
([#998](https://github.com/open-telemetry/opentelemetry-python/pull/998))
11+
- Samplers to accept parent Context
12+
([#1267](https://github.com/open-telemetry/opentelemetry-python/pull/1267))
1113

1214
## Version 0.14b0
1315

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,7 @@ def start_span( # pylint: disable=too-many-locals
757757
"parent_span_context must be a SpanContext or None."
758758
)
759759

760+
# is_valid determines root span
760761
if parent_span_context is None or not parent_span_context.is_valid:
761762
parent_span_context = None
762763
trace_id = self.source.ids_generator.generate_trace_id()
@@ -773,7 +774,7 @@ def start_span( # pylint: disable=too-many-locals
773774
# The sampler may also add attributes to the newly-created span, e.g.
774775
# to include information about the sampling result.
775776
sampling_result = self.source.sampler.should_sample(
776-
parent_span_context, trace_id, name, attributes, links,
777+
context, trace_id, name, attributes, links,
777778
)
778779

779780
trace_flags = (

‎opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@
6666
from typing import Optional, Sequence
6767

6868
# pylint: disable=unused-import
69-
from opentelemetry.trace import Link, SpanContext
69+
from opentelemetry.context import Context
70+
from opentelemetry.trace import Link, get_current_span
7071
from opentelemetry.util.types import Attributes
7172

7273

@@ -113,11 +114,11 @@ class Sampler(abc.ABC):
113114
@abc.abstractmethod
114115
def should_sample(
115116
self,
116-
parent_span_context: Optional["SpanContext"],
117+
parent_context: Optional["Context"],
117118
trace_id: int,
118119
name: str,
119120
attributes: Attributes = None,
120-
links: Sequence["Link"] = (),
121+
links: Sequence["Link"] = None,
121122
) -> "SamplingResult":
122123
pass
123124

@@ -134,11 +135,11 @@ def __init__(self, decision: "Decision"):
134135

135136
def should_sample(
136137
self,
137-
parent_span_context: Optional["SpanContext"],
138+
parent_context: Optional["Context"],
138139
trace_id: int,
139140
name: str,
140141
attributes: Attributes = None,
141-
links: Sequence["Link"] = (),
142+
links: Sequence["Link"] = None,
142143
) -> "SamplingResult":
143144
if self._decision is Decision.DROP:
144145
return SamplingResult(self._decision)
@@ -188,11 +189,11 @@ def bound(self) -> int:
188189

189190
def should_sample(
190191
self,
191-
parent_span_context: Optional["SpanContext"],
192+
parent_context: Optional["Context"],
192193
trace_id: int,
193194
name: str,
194-
attributes: Attributes = None, # TODO
195-
links: Sequence["Link"] = (),
195+
attributes: Attributes = None,
196+
links: Sequence["Link"] = None,
196197
) -> "SamplingResult":
197198
decision = Decision.DROP
198199
if trace_id & self.TRACE_ID_LIMIT < self.bound:
@@ -220,22 +221,27 @@ def __init__(self, delegate: Sampler):
220221

221222
def should_sample(
222223
self,
223-
parent_span_context: Optional["SpanContext"],
224+
parent_context: Optional["Context"],
224225
trace_id: int,
225226
name: str,
226-
attributes: Attributes = None, # TODO
227-
links: Sequence["Link"] = (),
227+
attributes: Attributes = None,
228+
links: Sequence["Link"] = None,
228229
) -> "SamplingResult":
229-
if parent_span_context is not None:
230+
if parent_context is not None:
231+
parent_span_context = get_current_span(
232+
parent_context
233+
).get_span_context()
234+
# only drop if parent exists and is not a root span
230235
if (
231-
not parent_span_context.is_valid
232-
or not parent_span_context.trace_flags.sampled
236+
parent_span_context is not None
237+
and parent_span_context.is_valid
238+
and not parent_span_context.trace_flags.sampled
233239
):
234240
return SamplingResult(Decision.DROP)
235241
return SamplingResult(Decision.RECORD_AND_SAMPLE, attributes)
236242

237243
return self._delegate.should_sample(
238-
parent_span_context=parent_span_context,
244+
parent_context=parent_context,
239245
trace_id=trace_id,
240246
name=name,
241247
attributes=attributes,

‎opentelemetry-sdk/tests/trace/test_sampling.py

Lines changed: 71 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -109,24 +109,34 @@ def test_always_off(self):
109109
self.assertEqual(sampled_always_on.attributes, {})
110110

111111
def test_default_on(self):
112+
context = trace.set_span_in_context(
113+
trace.DefaultSpan(
114+
trace.SpanContext(
115+
0xDEADBEEF,
116+
0xDEADBEF0,
117+
is_remote=False,
118+
trace_flags=TO_DEFAULT,
119+
)
120+
)
121+
)
112122
no_record_default_on = sampling.DEFAULT_ON.should_sample(
113-
trace.SpanContext(
114-
0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_DEFAULT
115-
),
116-
0xDEADBEF1,
117-
0xDEADBEF2,
118-
"unsampled parent, sampling on",
123+
context, 0xDEADBEF1, 0xDEADBEF2, "unsampled parent, sampling on",
119124
)
120125
self.assertFalse(no_record_default_on.decision.is_sampled())
121126
self.assertEqual(no_record_default_on.attributes, {})
122127

128+
context = trace.set_span_in_context(
129+
trace.DefaultSpan(
130+
trace.SpanContext(
131+
0xDEADBEEF,
132+
0xDEADBEF0,
133+
is_remote=False,
134+
trace_flags=TO_SAMPLED,
135+
)
136+
)
137+
)
123138
sampled_default_on = sampling.DEFAULT_ON.should_sample(
124-
trace.SpanContext(
125-
0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_SAMPLED
126-
),
127-
0xDEADBEF1,
128-
0xDEADBEF2,
129-
{"sampled parent": "sampling on"},
139+
context, 0xDEADBEF1, 0xDEADBEF2, {"sampled parent": "sampling on"},
130140
)
131141
self.assertTrue(sampled_default_on.decision.is_sampled())
132142
self.assertEqual(
@@ -142,24 +152,34 @@ def test_default_on(self):
142152
)
143153

144154
def test_default_off(self):
155+
context = trace.set_span_in_context(
156+
trace.DefaultSpan(
157+
trace.SpanContext(
158+
0xDEADBEEF,
159+
0xDEADBEF0,
160+
is_remote=False,
161+
trace_flags=TO_DEFAULT,
162+
)
163+
)
164+
)
145165
no_record_default_off = sampling.DEFAULT_OFF.should_sample(
146-
trace.SpanContext(
147-
0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_DEFAULT
148-
),
149-
0xDEADBEF1,
150-
0xDEADBEF2,
151-
"unsampled parent, sampling off",
166+
context, 0xDEADBEF1, 0xDEADBEF2, "unsampled parent, sampling off",
152167
)
153168
self.assertFalse(no_record_default_off.decision.is_sampled())
154169
self.assertEqual(no_record_default_off.attributes, {})
155170

171+
context = trace.set_span_in_context(
172+
trace.DefaultSpan(
173+
trace.SpanContext(
174+
0xDEADBEEF,
175+
0xDEADBEF0,
176+
is_remote=False,
177+
trace_flags=TO_SAMPLED,
178+
)
179+
)
180+
)
156181
sampled_default_off = sampling.DEFAULT_OFF.should_sample(
157-
trace.SpanContext(
158-
0xDEADBEEF, 0xDEADBEF0, is_remote=False, trace_flags=TO_SAMPLED
159-
),
160-
0xDEADBEF1,
161-
0xDEADBEF2,
162-
{"sampled parent": "sampling on"},
182+
context, 0xDEADBEF1, 0xDEADBEF2, {"sampled parent": "sampling on"},
163183
)
164184
self.assertTrue(sampled_default_off.decision.is_sampled())
165185
self.assertEqual(
@@ -275,32 +295,45 @@ def test_probability_sampler_limits(self):
275295

276296
def test_parent_based(self):
277297
sampler = sampling.ParentBased(sampling.ALWAYS_ON)
278-
# Check that the sampling decision matches the parent context if given
279-
self.assertFalse(
280-
sampler.should_sample(
298+
context = trace.set_span_in_context(
299+
trace.DefaultSpan(
281300
trace.SpanContext(
282301
0xDEADBEF0,
283302
0xDEADBEF1,
284303
is_remote=False,
285304
trace_flags=TO_DEFAULT,
286-
),
287-
0x7FFFFFFFFFFFFFFF,
288-
0xDEADBEEF,
289-
"span name",
305+
)
306+
)
307+
)
308+
# Check that the sampling decision matches the parent context if given
309+
self.assertFalse(
310+
sampler.should_sample(
311+
context, 0x7FFFFFFFFFFFFFFF, 0xDEADBEEF, "span name",
290312
).decision.is_sampled()
291313
)
292314

293-
sampler2 = sampling.ParentBased(sampling.ALWAYS_OFF)
294-
self.assertTrue(
295-
sampler2.should_sample(
315+
context = trace.set_span_in_context(
316+
trace.DefaultSpan(
296317
trace.SpanContext(
297318
0xDEADBEF0,
298319
0xDEADBEF1,
299320
is_remote=False,
300321
trace_flags=TO_SAMPLED,
301-
),
302-
0x8000000000000000,
303-
0xDEADBEEF,
304-
"span name",
322+
)
323+
)
324+
)
325+
sampler2 = sampling.ParentBased(sampling.ALWAYS_OFF)
326+
self.assertTrue(
327+
sampler2.should_sample(
328+
context, 0x8000000000000000, 0xDEADBEEF, "span name",
329+
).decision.is_sampled()
330+
)
331+
332+
# root span always sampled for parentbased
333+
context = trace.set_span_in_context(trace.INVALID_SPAN)
334+
sampler3 = sampling.ParentBased(sampling.ALWAYS_OFF)
335+
self.assertTrue(
336+
sampler3.should_sample(
337+
context, 0x8000000000000000, 0xDEADBEEF, "span name",
305338
).decision.is_sampled()
306339
)

0 commit comments

Comments
 (0)