fix: Do not use Celery signals for scope pushing/popping

untitaker · untitaker · commit 778e2078a518 · 2019-01-29T11:31:12.000+01:00
diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py
@@ -2,11 +2,11 @@
 
 import sys
 
-from celery.signals import task_failure, task_prerun, task_postrun
 from celery.exceptions import SoftTimeLimitExceeded
 
 from sentry_sdk.hub import Hub
 from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
+from sentry_sdk._compat import reraise
 from sentry_sdk.integrations import Integration
 from sentry_sdk.integrations.logging import ignore_logger
 
@@ -16,33 +16,42 @@ class CeleryIntegration(Integration):
 
     @staticmethod
     def setup_once():
-        task_prerun.connect(_handle_task_prerun, weak=False)
-        task_postrun.connect(_handle_task_postrun, weak=False)
-        task_failure.connect(_process_failure_signal, weak=False)
+        import celery.app.trace as trace
+
+        old_build_tracer = trace.build_tracer
+
+        def sentry_build_tracer(name, task, *args, **kwargs):
+            task.__call__ = _wrap_task_call(task, task.__call__)
+            task.run = _wrap_task_call(task, task.run)
+            return old_build_tracer(name, task, *args, **kwargs)
+
+        trace.build_tracer = sentry_build_tracer
 
         # 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")
 
 
-def _process_failure_signal(sender, task_id, einfo, **kw):
-    # einfo from celery is not reliable
-    exc_info = sys.exc_info()
-
-    hub = Hub.current
-    integration = hub.get_integration(CeleryIntegration)
-    if integration is None:
-        return
+def _wrap_task_call(self, f):
+    def _inner(*args, **kwargs):
+        hub = Hub.current
+        if hub.get_integration(CeleryIntegration) is None:
+            return f(*args, **kwargs)
 
-    _capture_event(hub, exc_info)
+        with hub.configure_scope() as scope:
+            if scope._name == "celery":
+                return f(*args, **kwargs)
 
+        with hub.push_scope() as scope:
+            scope._name = "celery"
+            scope.add_event_processor(_make_event_processor(args, kwargs, self))
+            try:
+                return f(*args, **kwargs)
+            except Exception:
+                reraise(*_capture_exception(hub))
 
-def _handle_task_prerun(sender, task, args, kwargs, **_):
-    hub = Hub.current
-    if hub.get_integration(CeleryIntegration) is not None:
-        scope = hub.push_scope().__enter__()
-        scope.add_event_processor(_make_event_processor(args, kwargs, task))
+    return _inner
 
 
 def _make_event_processor(args, kwargs, task):
@@ -78,16 +87,12 @@ def event_processor(event, hint):
     return event_processor
 
 
-def _handle_task_postrun(sender, task_id, task, **kw):
-    hub = Hub.current
-    if hub.get_integration(CeleryIntegration) is not None:
-        hub.pop_scope_unsafe()
-
-
-def _capture_event(hub, exc_info):
+def _capture_exception(hub):
+    exc_info = sys.exc_info()
     event, hint = event_from_exception(
         exc_info,
         client_options=hub.client.options,
         mechanism={"type": "celery", "handled": False},
     )
     hub.capture_event(event, hint=hint)
+    return exc_info
diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py
@@ -2,17 +2,35 @@
 
 pytest.importorskip("celery")
 
+from sentry_sdk import Hub, configure_scope
 from sentry_sdk.integrations.celery import CeleryIntegration
 
-from celery import Celery
+from celery import Celery, VERSION
 
 
 @pytest.fixture
-def celery(sentry_init):
-    sentry_init(integrations=[CeleryIntegration()])
-    celery = Celery(__name__)
-    celery.conf.CELERY_ALWAYS_EAGER = True
-    return celery
+def connect_signal(request):
+    def inner(signal, f):
+        signal.connect(f)
+        request.addfinalizer(lambda: signal.disconnect(f))
+
+    return inner
+
+
+@pytest.fixture
+def init_celery(sentry_init):
+    def inner():
+        sentry_init(integrations=[CeleryIntegration()])
+        celery = Celery(__name__)
+        celery.conf.CELERY_ALWAYS_EAGER = True
+        return celery
+
+    return inner
+
+
+@pytest.fixture
+def celery(init_celery):
+    return init_celery()
 
 
 def test_simple(capture_events, celery):
@@ -36,7 +54,7 @@ def dummy_task(x, y):
     exception, = event["exception"]["values"]
     assert exception["type"] == "ZeroDivisionError"
     assert exception["mechanism"]["type"] == "celery"
-    assert exception["stacktrace"]["frames"][1]["vars"]["foo"] == "42"
+    assert exception["stacktrace"]["frames"][0]["vars"]["foo"] == "42"
 
 
 def test_ignore_expected(capture_events, celery):
@@ -49,3 +67,32 @@ def dummy_task(x, y):
     dummy_task.delay(1, 2)
     dummy_task.delay(1, 0)
     assert not events
+
+
+def test_broken_prerun(capture_events, init_celery, connect_signal):
+    from celery.signals import task_prerun, task_postrun
+
+    def crash(*args, **kwargs):
+        1 / 0
+
+    # Order here is important to reproduce the bug: In Celery 3, a crashing
+    # prerun would prevent other preruns from running.
+
+    connect_signal(task_prerun, crash)
+    celery = init_celery()
+
+    events = capture_events()
+    assert len(Hub.current._stack) == 1
+
+    @celery.task(name="dummy_task")
+    def dummy_task(x, y):
+        assert len(Hub.current._stack) == 2
+        return x / y
+
+    try:
+        result = dummy_task.delay(2, 2)
+    except ZeroDivisionError:
+        if VERSION >= (4,):
+            raise
+
+    assert len(Hub.current._stack) == 1
