fix(celery): Fix dropped transactions under Celery 4.2+ (#825)

* Work around celery/celery#4875 which causes us to lose transaction events. Fix #824 
* Rewrite celery testsuite to use redis backend and test transactions too. This is better because it works on more celery versions (memory backend is often broken). However, this still does not trigger the bug, so I guess for this to be properly tested we'd need to install rabbitmq into CI? No thanks

Author: untitaker

.travis.yml

@@ -4,6 +4,7 @@ dist: xenial
 
 services:
   - postgresql
+  - redis-server
 
 language: python
 

sentry_sdk/integrations/celery.py

@@ -93,15 +93,23 @@ 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):
+                with capture_internal_exceptions():
+                    headers = dict(hub.iter_trace_propagation_headers())
+                    if headers:
+                        kwarg_headers = kwargs.setdefault("headers", {})
+                        kwarg_headers.update(headers)
+
+                        # https://github.com/celery/celery/issues/4875
+                        #
+                        # Need to setdefault the inner headers too since other
+                        # tracing tools (dd-trace-py) also employ this exact
+                        # workaround and we don't want to break them.
+                        #
+                        # This is not reproducible outside of AMQP, therefore no
+                        # tests!
+                        kwarg_headers.setdefault("headers", {}).update(headers)
+
                 return f(*args, **kwargs)
         else:
             return f(*args, **kwargs)
@@ -130,19 +138,22 @@ def _inner(*args, **kwargs):
             scope.clear_breadcrumbs()
             scope.add_event_processor(_make_event_processor(task, *args, **kwargs))
 
-            transaction = Transaction.continue_from_headers(
-                args[3].get("headers") or {},
-                op="celery.task",
-                name="unknown celery task",
-            )
-
-            # Could possibly use a better hook than this one
-            transaction.set_status("ok")
+            transaction = None
 
+            # Celery task objects are not a thing to be trusted. Even
+            # something such as attribute access can fail.
             with capture_internal_exceptions():
-                # Celery task objects are not a thing to be trusted. Even
-                # something such as attribute access can fail.
+                transaction = Transaction.continue_from_headers(
+                    args[3].get("headers") or {},
+                    op="celery.task",
+                    name="unknown celery task",
+                )
+
                 transaction.name = task.name
+                transaction.set_status("ok")
+
+            if transaction is None:
+                return f(*args, **kwargs)
 
             with hub.start_transaction(transaction):
                 return f(*args, **kwargs)

sentry_sdk/tracing.py

@@ -318,7 +318,7 @@ def set_status(self, value):
 
     def set_http_status(self, http_status):
         # type: (int) -> None
-        self.set_tag("http.status_code", http_status)
+        self.set_tag("http.status_code", str(http_status))
 
         if http_status < 400:
             self.set_status("ok")

tests/conftest.py

@@ -197,7 +197,7 @@ def append_event(event):
         def append_envelope(envelope):
             for item in envelope:
                 if item.headers.get("type") in ("event", "transaction"):
-                    events.append(item.payload.json)
+                    test_client.transport.capture_event(item.payload.json)
             return old_capture_envelope(envelope)
 
         monkeypatch.setattr(test_client.transport, "capture_event", append_event)
@@ -233,8 +233,14 @@ def append_envelope(envelope):
 
 
 @pytest.fixture
-def capture_events_forksafe(monkeypatch):
+def capture_events_forksafe(monkeypatch, capture_events, request):
     def inner():
+        in_process_events = capture_events()
+
+        @request.addfinalizer
+        def _():
+            assert not in_process_events
+
         events_r, events_w = os.pipe()
         events_r = os.fdopen(events_r, "rb", 0)
         events_w = os.fdopen(events_w, "wb", 0)

tests/integrations/celery/test_celery.py

@@ -22,17 +22,41 @@ def inner(signal, f):
 
 
 @pytest.fixture
-def init_celery(sentry_init):
-    def inner(propagate_traces=True, **kwargs):
+def init_celery(sentry_init, request):
+    def inner(propagate_traces=True, backend="always_eager", **kwargs):
         sentry_init(
             integrations=[CeleryIntegration(propagate_traces=propagate_traces)],
             **kwargs
         )
         celery = Celery(__name__)
-        if VERSION < (4,):
-            celery.conf.CELERY_ALWAYS_EAGER = True
+
+        if backend == "always_eager":
+            if VERSION < (4,):
+                celery.conf.CELERY_ALWAYS_EAGER = True
+            else:
+                celery.conf.task_always_eager = True
+        elif backend == "redis":
+            # broken on celery 3
+            if VERSION < (4,):
+                pytest.skip("Redis backend broken for some reason")
+
+            # this backend requires capture_events_forksafe
+            celery.conf.worker_max_tasks_per_child = 1
+            celery.conf.broker_url = "redis://127.0.0.1:6379"
+            celery.conf.result_backend = "redis://127.0.0.1:6379"
+            celery.conf.task_always_eager = False
+
+            Hub.main.bind_client(Hub.current.client)
+            request.addfinalizer(lambda: Hub.main.bind_client(None))
+
+            # Once we drop celery 3 we can use the celery_worker fixture
+            w = worker.worker(app=celery)
+            t = threading.Thread(target=w.run)
+            t.daemon = True
+            t.start()
         else:
-            celery.conf.task_always_eager = True
+            raise ValueError(backend)
+
         return celery
 
     return inner
@@ -273,15 +297,10 @@ def dummy_task(self):
 
 
 @pytest.mark.forked
-@pytest.mark.skipif(VERSION < (4,), reason="in-memory backend broken")
-def test_transport_shutdown(request, celery, capture_events_forksafe, tmpdir):
-    events = capture_events_forksafe()
+def test_redis_backend(init_celery, capture_events_forksafe, tmpdir):
+    celery = init_celery(traces_sample_rate=1.0, backend="redis", debug=True)
 
-    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
+    events = capture_events_forksafe()
 
     runs = []
 
@@ -290,21 +309,26 @@ 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()
+    # Curious: Cannot use delay() here or py2.7-celery-4.2 crashes
+    res = dummy_task.apply_async()
 
     with pytest.raises(Exception):
         # Celery 4.1 raises a gibberish exception
         res.wait()
 
+    # if this is nonempty, the worker never really forked
+    assert not runs
+
     event = events.read_event()
     (exception,) = event["exception"]["values"]
     assert exception["type"] == "ZeroDivisionError"
 
+    transaction = events.read_event()
+    assert (
+        transaction["contexts"]["trace"]["trace_id"]
+        == event["contexts"]["trace"]["trace_id"]
+    )
+
     events.read_flush()
 
     # if this is nonempty, the worker never really forked

tox.ini

@@ -38,7 +38,8 @@ envlist =
     {py3.6,py3.7}-sanic-19
 
     # TODO: Add py3.9
-    {pypy,py2.7,py3.5,py3.6,py3.7,py3.8}-celery-{4.1,4.2,4.3,4.4}
+    {pypy,py2.7,py3.5,py3.6}-celery-{4.1,4.2}
+    {pypy,py2.7,py3.5,py3.6,py3.7,py3.8}-celery-{4.3,4.4}
     {pypy,py2.7}-celery-3
 
     {py2.7,py3.7}-beam-{2.12,2.13}
@@ -128,6 +129,7 @@ deps =
     beam-2.13: apache-beam>=2.13.0, <2.14.0
     beam-master: git+https://github.com/apache/beam#egg=apache-beam&subdirectory=sdks/python
 
+    celery: redis
     celery-3: Celery>=3.1,<4.0
     celery-4.1: Celery>=4.1,<4.2
     celery-4.2: Celery>=4.2,<4.3