fix(celery): Don't set span status to error on control flow exception (#7912)

Author: sentrivana

sentry_sdk/integrations/celery/__init__.py

@@ -19,6 +19,7 @@
 from sentry_sdk.traces import BAGGAGE_HEADER_NAME, SegmentNameSource, Span
 from sentry_sdk.tracing_utils import Baggage
 from sentry_sdk.utils import (
+    _register_control_flow_exception,
     capture_internal_exceptions,
     event_from_exception,
     parse_version,
@@ -78,6 +79,7 @@ def setup_once() -> None:
         _patch_celery_send_task()
         _patch_worker_exit()
         _patch_producer_publish()
+        _register_control_flow_exception(list(CELERY_CONTROL_FLOW_EXCEPTIONS))
 
         # 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
@@ -90,25 +92,18 @@ def setup_once() -> None:
         ignore_logger_for_events("celery.redirected")
 
 
-def _set_status(status: str) -> None:
-    with capture_internal_exceptions():
-        span = sentry_sdk.get_current_span()
-
-        if span is not None:
-            span.status = "ok" if status == "ok" else "error"
-
-
 def _capture_exception(task: "Any", exc_info: "ExcInfo") -> None:
     client = sentry_sdk.get_client()
     if client.get_integration(CeleryIntegration) is None:
         return
 
     if isinstance(exc_info[1], CELERY_CONTROL_FLOW_EXCEPTIONS):
-        # ??? Doesn't map to anything
-        _set_status("aborted")
+        # Expected control flow exits, not errors
         return
 
-    _set_status("internal_error")
+    span = sentry_sdk.get_current_span()
+    if span is not None:
+        span.status = "error"
 
     if hasattr(task, "throws") and isinstance(exc_info[1], task.throws):
         return

tests/integrations/celery/test_celery.py

@@ -6,6 +6,7 @@
 import pytest
 from celery import VERSION, Celery
 from celery.bin import worker
+from celery.exceptions import Ignore, Reject, Retry
 
 import sentry_sdk
 from sentry_sdk.integrations.celery import (
@@ -409,6 +410,39 @@ def dummy_task(self):
         assert e["type"] == "ZeroDivisionError"
 
 
+@pytest.mark.parametrize(
+    "exception",
+    [Retry, Ignore, Reject],
+    ids=["retry", "ignore", "reject"],
+)
+def test_control_flow_exceptions_not_errors(init_celery, capture_items, exception):
+    celery = init_celery(traces_sample_rate=1.0)
+    should_raise = [False]
+
+    @celery.task(name="dummy_task")
+    def dummy_task():
+        if should_raise[0]:
+            raise exception()
+
+    # XXX: For some reason the first call does not get instrumented properly.
+    dummy_task.delay()
+    sentry_sdk.flush()
+
+    items = capture_items("event", "span")
+    should_raise[0] = True
+
+    dummy_task.delay()
+    sentry_sdk.flush()
+
+    assert not [item for item in items if item.type == "event"]
+
+    process_span, execution_span = [item.payload for item in items]
+    assert process_span["attributes"]["sentry.op"] == "queue.process"
+    assert process_span["status"] == "ok"
+    assert execution_span["is_segment"] is True
+    assert execution_span["status"] == "ok"
+
+
 @pytest.mark.forked
 @pytest.mark.parametrize("newrelic_order", ["sentry_first", "sentry_last"])
 def test_newrelic_interference(init_celery, newrelic_order, celery_invocation):