Skip to content

Commit a843018

Browse files
mitsuhikountitaker
mitsuhiko
authored and
untitaker
committed
feat(tracing): Initial tracing experiments (getsentry#342)
* feat(tracing): Initial tracing experiments * ref: Hide traceparent header in public api * ref: Moved tracing code to wsgi integration * ref: lint * feat: Added basic celery propagation * feat: Added a way to disable trace propagation for celery * fix: Optional headers * ref: Always populate traces if enabled * ref: traceparent -> sentry-trace * fix: Linters
1 parent 05bfb9d commit a843018

File tree

9 files changed

+191
-5
lines changed

9 files changed

+191
-5
lines changed

‎sentry_sdk/consts.py

+2
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
"debug": bool,
3838
"attach_stacktrace": bool,
3939
"ca_certs": Optional[str],
40+
"propagate_traces": bool,
4041
},
4142
total=False,
4243
)
@@ -69,6 +70,7 @@
6970
"debug": False,
7071
"attach_stacktrace": False,
7172
"ca_certs": None,
73+
"propagate_traces": True,
7274
}
7375

7476

‎sentry_sdk/hub.py

+12
Original file line numberDiff line numberDiff line change
@@ -415,6 +415,18 @@ def flush(self, timeout=None, callback=None):
415415
if client is not None:
416416
return client.flush(timeout=timeout, callback=callback)
417417

418+
def iter_trace_propagation_headers(self):
419+
client, scope = self._stack[-1]
420+
if scope._span is None:
421+
return
422+
423+
propagate_traces = client and client.options["propagate_traces"]
424+
if not propagate_traces:
425+
return
426+
427+
for item in scope._span.iter_headers():
428+
yield item
429+
418430

419431
GLOBAL_HUB = Hub()
420432
_local.set(GLOBAL_HUB)

‎sentry_sdk/integrations/celery.py

+31
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from sentry_sdk.hub import Hub
88
from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
9+
from sentry_sdk.tracing import SpanContext
910
from sentry_sdk._compat import reraise
1011
from sentry_sdk.integrations import Integration
1112
from sentry_sdk.integrations.logging import ignore_logger
@@ -14,6 +15,9 @@
1415
class CeleryIntegration(Integration):
1516
identifier = "celery"
1617

18+
def __init__(self, propagate_traces=True):
19+
self.propagate_traces = propagate_traces
20+
1721
@staticmethod
1822
def setup_once():
1923
import celery.app.trace as trace # type: ignore
@@ -25,6 +29,7 @@ def sentry_build_tracer(name, task, *args, **kwargs):
2529
# short-circuits to task.run if it thinks it's safe.
2630
task.__call__ = _wrap_task_call(task, task.__call__)
2731
task.run = _wrap_task_call(task, task.run)
32+
task.apply_async = _wrap_apply_async(task, task.apply_async)
2833
return _wrap_tracer(task, old_build_tracer(name, task, *args, **kwargs))
2934

3035
trace.build_tracer = sentry_build_tracer
@@ -37,6 +42,23 @@ def sentry_build_tracer(name, task, *args, **kwargs):
3742
ignore_logger("celery.worker.job")
3843

3944

45+
def _wrap_apply_async(task, f):
46+
def apply_async(self, *args, **kwargs):
47+
hub = Hub.current
48+
integration = hub.get_integration(CeleryIntegration)
49+
if integration is not None and integration.propagate_traces:
50+
headers = None
51+
for key, value in hub.iter_trace_propagation_headers():
52+
if headers is None:
53+
headers = dict(kwargs.get("headers") or {})
54+
headers[key] = value
55+
if headers is not None:
56+
kwargs["headers"] = headers
57+
return f(self, *args, **kwargs)
58+
59+
return apply_async
60+
61+
4062
def _wrap_tracer(task, f):
4163
# Need to wrap tracer for pushing the scope before prerun is sent, and
4264
# popping it after postrun is sent.
@@ -52,13 +74,22 @@ def _inner(*args, **kwargs):
5274
with hub.push_scope() as scope:
5375
scope._name = "celery"
5476
scope.clear_breadcrumbs()
77+
_continue_trace(args[3].get("headers") or {}, scope)
5578
scope.add_event_processor(_make_event_processor(task, *args, **kwargs))
5679

5780
return f(*args, **kwargs)
5881

5982
return _inner
6083

6184

85+
def _continue_trace(headers, scope):
86+
if headers:
87+
span_context = SpanContext.continue_from_headers(headers)
88+
else:
89+
span_context = SpanContext.start_trace()
90+
scope.set_span_context(span_context)
91+
92+
6293
def _wrap_task_call(task, f):
6394
# Need to wrap task call because the exception is caught before we get to
6495
# see it. Also celery's reported stacktrace is untrustworthy.

‎sentry_sdk/integrations/flask.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -96,9 +96,10 @@ def _request_started(sender, **kwargs):
9696
if integration is None:
9797
return
9898

99-
weak_request = weakref.ref(_request_ctx_stack.top.request)
10099
app = _app_ctx_stack.top.app
101100
with hub.configure_scope() as scope:
101+
request = _request_ctx_stack.top.request
102+
weak_request = weakref.ref(request)
102103
scope.add_event_processor(
103104
_make_request_event_processor( # type: ignore
104105
app, weak_request, integration

‎sentry_sdk/integrations/stdlib.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,8 @@ def install_httplib():
2424

2525
def putrequest(self, method, url, *args, **kwargs):
2626
rv = real_putrequest(self, method, url, *args, **kwargs)
27-
if Hub.current.get_integration(StdlibIntegration) is None:
27+
hub = Hub.current
28+
if hub.get_integration(StdlibIntegration) is None:
2829
return rv
2930

3031
self._sentrysdk_data_dict = data = {}
@@ -42,6 +43,9 @@ def putrequest(self, method, url, *args, **kwargs):
4243
url,
4344
)
4445

46+
for key, value in hub.iter_trace_propagation_headers():
47+
self.putheader(key, value)
48+
4549
data["url"] = real_url
4650
data["method"] = method
4751
return rv

‎sentry_sdk/integrations/wsgi.py

+2
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from sentry_sdk.hub import Hub, _should_send_default_pii
44
from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
55
from sentry_sdk._compat import PY2, reraise
6+
from sentry_sdk.tracing import SpanContext
67
from sentry_sdk.integrations._wsgi_common import _filter_headers
78

89
if False:
@@ -81,6 +82,7 @@ def __call__(self, environ, start_response):
8182
with hub.configure_scope() as scope:
8283
scope.clear_breadcrumbs()
8384
scope._name = "wsgi"
85+
scope.set_span_context(SpanContext.continue_from_environ(environ))
8486
scope.add_event_processor(_make_wsgi_event_processor(environ))
8587

8688
try:

‎sentry_sdk/scope.py

+14
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ class Scope(object):
5959
"_event_processors",
6060
"_error_processors",
6161
"_should_capture",
62+
"_span",
6263
)
6364

6465
def __init__(self):
@@ -88,6 +89,10 @@ def user(self, value):
8889
"""When set a specific user is bound to the scope."""
8990
self._user = value
9091

92+
def set_span_context(self, span_context):
93+
"""Sets the span context."""
94+
self._span = span_context
95+
9196
def set_tag(self, key, value):
9297
"""Sets a tag for a key to a specific value."""
9398
self._tags[key] = value
@@ -127,6 +132,8 @@ def clear(self):
127132
self.clear_breadcrumbs()
128133
self._should_capture = True
129134

135+
self._span = None
136+
130137
def clear_breadcrumbs(self):
131138
# type: () -> None
132139
"""Clears breadcrumb buffer."""
@@ -193,6 +200,12 @@ def _drop(event, cause, ty):
193200
if self._contexts:
194201
event.setdefault("contexts", {}).update(self._contexts)
195202

203+
if self._span is not None:
204+
event.setdefault("contexts", {})["trace"] = {
205+
"trace_id": self._span.trace_id,
206+
"span_id": self._span.span_id,
207+
}
208+
196209
exc_info = hint.get("exc_info") if hint is not None else None
197210
if exc_info is not None:
198211
for processor in self._error_processors:
@@ -230,6 +243,7 @@ def __copy__(self):
230243
rv._error_processors = list(self._error_processors)
231244

232245
rv._should_capture = self._should_capture
246+
rv._span = self._span
233247

234248
return rv
235249

‎sentry_sdk/tracing.py

+93
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import re
2+
import uuid
3+
4+
_traceparent_header_format_re = re.compile(
5+
"^[ \t]*([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})" "(-.*)?[ \t]*$"
6+
)
7+
8+
9+
class _EnvironHeaders(object):
10+
def __init__(self, environ):
11+
self.environ = environ
12+
13+
def get(self, key):
14+
return self.environ.get("HTTP_" + key.replace("-", "_").upper())
15+
16+
17+
class SpanContext(object):
18+
def __init__(self, trace_id, span_id, recorded=False, parent=None):
19+
self.trace_id = trace_id
20+
self.span_id = span_id
21+
self.recorded = recorded
22+
self.parent = None
23+
24+
def __repr__(self):
25+
return "%s(trace_id=%r, span_id=%r, recorded=%r)" % (
26+
self.__class__.__name__,
27+
self.trace_id,
28+
self.span_id,
29+
self.recorded,
30+
)
31+
32+
@classmethod
33+
def start_trace(cls, recorded=False):
34+
return cls(
35+
trace_id=uuid.uuid4().hex, span_id=uuid.uuid4().hex[16:], recorded=recorded
36+
)
37+
38+
def new_span(self):
39+
if self.trace_id is None:
40+
return SpanContext.start_trace()
41+
return SpanContext(
42+
trace_id=self.trace_id,
43+
span_id=uuid.uuid4().hex[16:],
44+
parent=self,
45+
recorded=self.recorded,
46+
)
47+
48+
@classmethod
49+
def continue_from_environ(cls, environ):
50+
return cls.continue_from_headers(_EnvironHeaders(environ))
51+
52+
@classmethod
53+
def continue_from_headers(cls, headers):
54+
parent = cls.from_traceparent(headers.get("sentry-trace"))
55+
if parent is None:
56+
return cls.start_trace()
57+
return parent.new_span()
58+
59+
def iter_headers(self):
60+
yield "sentry-trace", self.to_traceparent()
61+
62+
@classmethod
63+
def from_traceparent(cls, traceparent):
64+
if not traceparent:
65+
return None
66+
67+
match = _traceparent_header_format_re.match(traceparent)
68+
if match is None:
69+
return None
70+
71+
version, trace_id, span_id, trace_options, extra = match.groups()
72+
73+
if int(trace_id, 16) == 0 or int(span_id, 16) == 0:
74+
return None
75+
76+
version = int(version, 16)
77+
if version == 0:
78+
if extra:
79+
return None
80+
elif version == 255:
81+
return None
82+
83+
options = int(trace_options, 16)
84+
85+
return cls(trace_id=trace_id, span_id=span_id, recorded=options & 1 != 0)
86+
87+
def to_traceparent(self):
88+
return "%02x-%s-%s-%02x" % (
89+
0,
90+
self.trace_id,
91+
self.span_id,
92+
self.recorded and 1 or 0,
93+
)

‎tests/integrations/celery/test_celery.py

+30-3
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@
44

55
pytest.importorskip("celery")
66

7-
from sentry_sdk import Hub
7+
from sentry_sdk import Hub, configure_scope
88
from sentry_sdk.integrations.celery import CeleryIntegration
9+
from sentry_sdk.tracing import SpanContext
910

1011
from celery import Celery, VERSION
1112
from celery.bin import worker
@@ -22,8 +23,8 @@ def inner(signal, f):
2223

2324
@pytest.fixture
2425
def init_celery(sentry_init):
25-
def inner():
26-
sentry_init(integrations=[CeleryIntegration()])
26+
def inner(propagate_traces=True):
27+
sentry_init(integrations=[CeleryIntegration(propagate_traces=propagate_traces)])
2728
celery = Celery(__name__)
2829
if VERSION < (4,):
2930
celery.conf.CELERY_ALWAYS_EAGER = True
@@ -47,9 +48,15 @@ def dummy_task(x, y):
4748
foo = 42 # noqa
4849
return x / y
4950

51+
span_context = SpanContext.start_trace()
52+
with configure_scope() as scope:
53+
scope.set_span_context(span_context)
5054
dummy_task.delay(1, 2)
5155
dummy_task.delay(1, 0)
56+
5257
event, = events
58+
assert event["contexts"]["trace"]["trace_id"] == span_context.trace_id
59+
assert event["contexts"]["trace"]["span_id"] != span_context.span_id
5360
assert event["transaction"] == "dummy_task"
5461
assert event["extra"]["celery-job"] == {
5562
"args": [1, 0],
@@ -63,6 +70,26 @@ def dummy_task(x, y):
6370
assert exception["stacktrace"]["frames"][0]["vars"]["foo"] == "42"
6471

6572

73+
def test_simple_no_propagation(capture_events, init_celery):
74+
celery = init_celery(propagate_traces=False)
75+
events = capture_events()
76+
77+
@celery.task(name="dummy_task")
78+
def dummy_task():
79+
1 / 0
80+
81+
span_context = SpanContext.start_trace()
82+
with configure_scope() as scope:
83+
scope.set_span_context(span_context)
84+
dummy_task.delay()
85+
86+
event, = events
87+
assert event["contexts"]["trace"]["trace_id"] != span_context.trace_id
88+
assert event["transaction"] == "dummy_task"
89+
exception, = event["exception"]["values"]
90+
assert exception["type"] == "ZeroDivisionError"
91+
92+
6693
def test_ignore_expected(capture_events, celery):
6794
events = capture_events()
6895

0 commit comments

Comments
 (0)