fix(celery): flush queue on worker exit (#295)

* fix(celery): flush queue on worker exit

Fix #285

* fix: Fix tests

* fix: Linters

Author: untitaker

sentry_sdk/integrations/celery.py

@@ -23,12 +23,14 @@ def setup_once():
         def sentry_build_tracer(name, task, *args, **kwargs):
             # 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.__call__)
-            task.run = _wrap_task_call(task.run)
+            task.__call__ = _wrap_task_call(task, task.__call__)
+            task.run = _wrap_task_call(task, task.run)
             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>".
@@ -57,14 +59,17 @@ def _inner(*args, **kwargs):
     return _inner
 
 
-def _wrap_task_call(f):
+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.
     def _inner(*args, **kwargs):
         try:
             return f(*args, **kwargs)
         except Exception:
-            reraise(*_capture_exception())
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _capture_exception(task, exc_info)
+            reraise(*exc_info)
 
     return _inner
 
@@ -83,15 +88,6 @@ def event_processor(event, hint):
             }
 
         if "exc_info" in hint:
-            with capture_internal_exceptions():
-                if isinstance(hint["exc_info"][1], Retry):
-                    return None
-
-                if hasattr(task, "throws") and isinstance(
-                    hint["exc_info"][1], task.throws
-                ):
-                    return None
-
             with capture_internal_exceptions():
                 if issubclass(hint["exc_info"][0], SoftTimeLimitExceeded):
                     event["fingerprint"] = [
@@ -105,16 +101,39 @@ def event_processor(event, hint):
     return event_processor
 
 
-def _capture_exception():
+def _capture_exception(task, exc_info):
     hub = Hub.current
-    exc_info = sys.exc_info()
 
-    if hub.get_integration(CeleryIntegration) is not None:
-        event, hint = event_from_exception(
-            exc_info,
-            client_options=hub.client.options,
-            mechanism={"type": "celery", "handled": False},
-        )
-        hub.capture_event(event, hint=hint)
+    if hub.get_integration(CeleryIntegration) is None:
+        return
+    if isinstance(exc_info[1], Retry):
+        return
+    if hasattr(task, "throws") and isinstance(exc_info[1], task.throws):
+        return
+
+    event, hint = event_from_exception(
+        exc_info,
+        client_options=hub.client.options,
+        mechanism={"type": "celery", "handled": False},
+    )
+
+    hub.capture_event(event, hint=hint)
+
+
+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()
 
-    return exc_info
+    Worker.workloop = sentry_workloop

tests/conftest.py

@@ -122,3 +122,41 @@ def append(event):
         return events
 
     return inner
+
+
+@pytest.fixture
+def capture_events_forksafe(monkeypatch):
+    def inner():
+        events_r, events_w = os.pipe()
+        events_r = os.fdopen(events_r, "rb", 0)
+        events_w = os.fdopen(events_w, "wb", 0)
+
+        test_client = sentry_sdk.Hub.current.client
+
+        old_capture_event = test_client.transport.capture_event
+
+        def append(event):
+            events_w.write(json.dumps(event).encode("utf-8"))
+            events_w.write(b"\n")
+            return old_capture_event(event)
+
+        def flush(timeout=None, callback=None):
+            events_w.write(b"flush\n")
+
+        monkeypatch.setattr(test_client.transport, "capture_event", append)
+        monkeypatch.setattr(test_client, "flush", flush)
+
+        return EventStreamReader(events_r)
+
+    return inner
+
+
+class EventStreamReader(object):
+    def __init__(self, file):
+        self.file = file
+
+    def read_event(self):
+        return json.loads(self.file.readline().decode("utf-8"))
+
+    def read_flush(self):
+        assert self.file.readline() == b"flush\n"

tests/integrations/celery/test_celery.py

@@ -1,3 +1,5 @@
+import threading
+
 import pytest
 
 pytest.importorskip("celery")
@@ -6,6 +8,7 @@
 from sentry_sdk.integrations.celery import CeleryIntegration
 
 from celery import Celery, VERSION
+from celery.bin import worker
 
 
 @pytest.fixture
@@ -22,7 +25,10 @@ def init_celery(sentry_init):
     def inner():
         sentry_init(integrations=[CeleryIntegration()])
         celery = Celery(__name__)
-        celery.conf.CELERY_ALWAYS_EAGER = True
+        if VERSION < (4,):
+            celery.conf.CELERY_ALWAYS_EAGER = True
+        else:
+            celery.conf.task_always_eager = True
         return celery
 
     return inner
@@ -92,11 +98,11 @@ def dummy_task(x, y):
         stack_lengths.append(len(Hub.current._stack))
         return x / y
 
-    try:
+    if VERSION >= (4,):
         dummy_task.delay(2, 2)
-    except ZeroDivisionError:
-        if VERSION >= (4,):
-            raise
+    else:
+        with pytest.raises(ZeroDivisionError):
+            dummy_task.delay(2, 2)
 
     assert len(Hub.current._stack) == 1
     if VERSION < (4,):
@@ -139,3 +145,41 @@ def dummy_task(self):
 
     for e in exceptions:
         assert e["type"] == "ZeroDivisionError"
+
+
+@pytest.mark.skipif(VERSION < (4,), reason="in-memory backend broken")
+def test_transport_shutdown(request, celery, capture_events_forksafe, tmpdir):
+    events = capture_events_forksafe()
+
+    celery.conf.worker_max_tasks_per_child = 1
+    celery.conf.broker_url = "memory://localhost/"
+    celery.conf.broker_backend = "memory"
+    celery.conf.result_backend = "file://{}".format(tmpdir.mkdir("celery-results"))
+    celery.conf.task_always_eager = False
+
+    runs = []
+
+    @celery.task(name="dummy_task", bind=True)
+    def dummy_task(self):
+        runs.append(1)
+        1 / 0
+
+    res = dummy_task.delay()
+
+    w = worker.worker(app=celery)
+    t = threading.Thread(target=w.run)
+    t.daemon = True
+    t.start()
+
+    with pytest.raises(Exception):
+        # Celery 4.1 raises a gibberish exception
+        res.wait()
+
+    event = events.read_event()
+    exception, = event["exception"]["values"]
+    assert exception["type"] == "ZeroDivisionError"
+
+    events.read_flush()
+
+    # if this is nonempty, the worker never really forked
+    assert not runs

tests/integrations/rq/test_rq.py

@@ -1,13 +1,8 @@
 from sentry_sdk.integrations.rq import RqIntegration
 
-import os
-import json
-
 from fakeredis import FakeStrictRedis
 import rq
 
-from sentry_sdk import Hub
-
 
 def crashing_job(foo):
     1 / 0
@@ -40,32 +35,19 @@ def test_basic(sentry_init, capture_events):
     }
 
 
-def test_transport_shutdown(sentry_init):
+def test_transport_shutdown(sentry_init, capture_events_forksafe):
     sentry_init(integrations=[RqIntegration()])
 
-    events_r, events_w = os.pipe()
-    events_r = os.fdopen(events_r, "rb", 0)
-    events_w = os.fdopen(events_w, "wb", 0)
-
-    def capture_event(event):
-        events_w.write(json.dumps(event).encode("utf-8"))
-        events_w.write(b"\n")
-
-    def flush(timeout=None, callback=None):
-        events_w.write(b"flush\n")
-
-    Hub.current.client.transport.capture_event = capture_event
-    Hub.current.client.flush = flush
+    events = capture_events_forksafe()
 
     queue = rq.Queue(connection=FakeStrictRedis())
     worker = rq.Worker([queue], connection=queue.connection)
 
     queue.enqueue(crashing_job, foo=42)
     worker.work(burst=True)
 
-    event = events_r.readline()
-    event = json.loads(event.decode("utf-8"))
+    event = events.read_event()
+    events.read_flush()
+
     exception, = event["exception"]["values"]
     assert exception["type"] == "ZeroDivisionError"
-
-    assert events_r.readline() == b"flush\n"