210 lines
6.7 KiB
Python
210 lines
6.7 KiB
Python
from __future__ import absolute_import
|
|
|
|
import functools
|
|
import sys
|
|
|
|
from celery.exceptions import ( # type: ignore
|
|
SoftTimeLimitExceeded,
|
|
Retry,
|
|
Ignore,
|
|
Reject,
|
|
)
|
|
|
|
from sentry_sdk.hub import Hub
|
|
from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
|
|
from sentry_sdk.tracing import Span
|
|
from sentry_sdk._compat import reraise
|
|
from sentry_sdk.integrations import Integration
|
|
from sentry_sdk.integrations.logging import ignore_logger
|
|
from sentry_sdk._types import MYPY
|
|
|
|
if MYPY:
|
|
from typing import Any
|
|
|
|
|
|
CELERY_CONTROL_FLOW_EXCEPTIONS = (Retry, Ignore, Reject)
|
|
|
|
|
|
class CeleryIntegration(Integration):
|
|
identifier = "celery"
|
|
|
|
def __init__(self, propagate_traces=True):
|
|
# type: (bool) -> None
|
|
self.propagate_traces = propagate_traces
|
|
|
|
@staticmethod
|
|
def setup_once():
|
|
# type: () -> None
|
|
import celery.app.trace as trace # type: ignore
|
|
|
|
old_build_tracer = trace.build_tracer
|
|
|
|
def sentry_build_tracer(name, task, *args, **kwargs):
|
|
if not getattr(task, "_sentry_is_patched", False):
|
|
# Need to patch both methods because older celery sometimes
|
|
# short-circuits to task.run if it thinks it's safe.
|
|
task.__call__ = _wrap_task_call(task, task.__call__)
|
|
task.run = _wrap_task_call(task, task.run)
|
|
task.apply_async = _wrap_apply_async(task, task.apply_async)
|
|
|
|
# `build_tracer` is apparently called for every task
|
|
# invocation. Can't wrap every celery task for every invocation
|
|
# or we will get infinitely nested wrapper functions.
|
|
task._sentry_is_patched = True
|
|
|
|
return _wrap_tracer(task, old_build_tracer(name, task, *args, **kwargs))
|
|
|
|
trace.build_tracer = sentry_build_tracer
|
|
|
|
_patch_worker_exit()
|
|
|
|
# This logger logs every status of every task that ran on the worker.
|
|
# Meaning that every task's breadcrumbs are full of stuff like "Task
|
|
# <foo> raised unexpected <bar>".
|
|
ignore_logger("celery.worker.job")
|
|
ignore_logger("celery.app.trace")
|
|
|
|
|
|
def _wrap_apply_async(task, f):
|
|
@functools.wraps(f)
|
|
def apply_async(*args, **kwargs):
|
|
hub = Hub.current
|
|
integration = hub.get_integration(CeleryIntegration)
|
|
if integration is not None and integration.propagate_traces:
|
|
headers = None
|
|
for key, value in hub.iter_trace_propagation_headers():
|
|
if headers is None:
|
|
headers = dict(kwargs.get("headers") or {})
|
|
headers[key] = value
|
|
if headers is not None:
|
|
kwargs["headers"] = headers
|
|
|
|
with hub.start_span(op="celery.submit", description=task.name):
|
|
return f(*args, **kwargs)
|
|
else:
|
|
return f(*args, **kwargs)
|
|
|
|
return apply_async
|
|
|
|
|
|
def _wrap_tracer(task, f):
|
|
# Need to wrap tracer for pushing the scope before prerun is sent, and
|
|
# popping it after postrun is sent.
|
|
#
|
|
# This is the reason we don't use signals for hooking in the first place.
|
|
# Also because in Celery 3, signal dispatch returns early if one handler
|
|
# crashes.
|
|
@functools.wraps(f)
|
|
def _inner(*args, **kwargs):
|
|
hub = Hub.current
|
|
if hub.get_integration(CeleryIntegration) is None:
|
|
return f(*args, **kwargs)
|
|
|
|
with hub.push_scope() as scope:
|
|
scope._name = "celery"
|
|
scope.clear_breadcrumbs()
|
|
scope.add_event_processor(_make_event_processor(task, *args, **kwargs))
|
|
|
|
span = Span.continue_from_headers(args[3].get("headers") or {})
|
|
span.op = "celery.task"
|
|
span.transaction = "unknown celery task"
|
|
|
|
with capture_internal_exceptions():
|
|
# Celery task objects are not a thing to be trusted. Even
|
|
# something such as attribute access can fail.
|
|
span.transaction = task.name
|
|
|
|
with hub.start_span(span):
|
|
return f(*args, **kwargs)
|
|
|
|
return _inner
|
|
|
|
|
|
def _wrap_task_call(task, f):
|
|
# Need to wrap task call because the exception is caught before we get to
|
|
# see it. Also celery's reported stacktrace is untrustworthy.
|
|
|
|
# functools.wraps is important here because celery-once looks at this
|
|
# method's name.
|
|
# https://github.com/getsentry/sentry-python/issues/421
|
|
@functools.wraps(f)
|
|
def _inner(*args, **kwargs):
|
|
try:
|
|
return f(*args, **kwargs)
|
|
except Exception:
|
|
exc_info = sys.exc_info()
|
|
with capture_internal_exceptions():
|
|
_capture_exception(task, exc_info)
|
|
reraise(*exc_info)
|
|
|
|
return _inner
|
|
|
|
|
|
def _make_event_processor(task, uuid, args, kwargs, request=None):
|
|
def event_processor(event, hint):
|
|
with capture_internal_exceptions():
|
|
extra = event.setdefault("extra", {})
|
|
extra["celery-job"] = {
|
|
"task_name": task.name,
|
|
"args": args,
|
|
"kwargs": kwargs,
|
|
}
|
|
|
|
if "exc_info" in hint:
|
|
with capture_internal_exceptions():
|
|
if issubclass(hint["exc_info"][0], SoftTimeLimitExceeded):
|
|
event["fingerprint"] = [
|
|
"celery",
|
|
"SoftTimeLimitExceeded",
|
|
getattr(task, "name", task),
|
|
]
|
|
|
|
return event
|
|
|
|
return event_processor
|
|
|
|
|
|
def _capture_exception(task, exc_info):
|
|
hub = Hub.current
|
|
|
|
if hub.get_integration(CeleryIntegration) is None:
|
|
return
|
|
if isinstance(exc_info[1], CELERY_CONTROL_FLOW_EXCEPTIONS):
|
|
return
|
|
if hasattr(task, "throws") and isinstance(exc_info[1], task.throws):
|
|
return
|
|
|
|
# If an integration is there, a client has to be there.
|
|
client = hub.client # type: Any
|
|
|
|
event, hint = event_from_exception(
|
|
exc_info,
|
|
client_options=client.options,
|
|
mechanism={"type": "celery", "handled": False},
|
|
)
|
|
|
|
hub.capture_event(event, hint=hint)
|
|
|
|
with capture_internal_exceptions():
|
|
with hub.configure_scope() as scope:
|
|
scope.span.set_failure()
|
|
|
|
|
|
def _patch_worker_exit():
|
|
# Need to flush queue before worker shutdown because a crashing worker will
|
|
# call os._exit
|
|
from billiard.pool import Worker # type: ignore
|
|
|
|
old_workloop = Worker.workloop
|
|
|
|
def sentry_workloop(*args, **kwargs):
|
|
try:
|
|
return old_workloop(*args, **kwargs)
|
|
finally:
|
|
with capture_internal_exceptions():
|
|
hub = Hub.current
|
|
if hub.get_integration(CeleryIntegration) is not None:
|
|
hub.flush()
|
|
|
|
Worker.workloop = sentry_workloop
|