Set the correct status on spans

Author: antonpirker

sentry_sdk/integrations/celery/__init__.py

@@ -113,7 +113,6 @@ def _capture_exception(task, exc_info):
         return
 
     if isinstance(exc_info[1], CELERY_CONTROL_FLOW_EXCEPTIONS):
-        # ??? Doesn't map to anything
         _set_status("aborted")
         return
 
@@ -277,16 +276,25 @@ def apply_async(*args, **kwargs):
                 op=OP.QUEUE_SUBMIT_CELERY,
                 name=task_name,
                 origin=CeleryIntegration.origin,
+                only_if_parent=True,
             )
             if not task_started_from_beat
             else NoOpMgr()
         )  # type: Union[Span, NoOpMgr]
 
         with span_mgr as span:
-            kwargs["headers"] = _update_celery_task_headers(
-                kwarg_headers, span, integration.monitor_beat_tasks
-            )
-            return f(*args, **kwargs)
+            try:
+                kwargs["headers"] = _update_celery_task_headers(
+                    kwarg_headers, span, integration.monitor_beat_tasks
+                )
+                return_value = f(*args, **kwargs)
+
+            except Exception:
+                reraise(*sys.exc_info())
+            else:
+                span.set_status(SPANSTATUS.OK)
+
+            return return_value
 
     return apply_async  # type: ignore
 
@@ -314,24 +322,30 @@ def _inner(*args, **kwargs):
             # something such as attribute access can fail.
             headers = args[3].get("headers") or {}
             with sentry_sdk.continue_trace(headers):
-                with sentry_sdk.start_span(
-                    op=OP.QUEUE_TASK_CELERY,
-                    name=task.name,
-                    source=TRANSACTION_SOURCE_TASK,
-                    origin=CeleryIntegration.origin,
-                    custom_sampling_context={
-                        "celery_job": {
-                            "task": task.name,
-                            # for some reason, args[1] is a list if non-empty but a
-                            # tuple if empty
-                            "args": list(args[1]),
-                            "kwargs": args[2],
-                        }
-                    },
-                ) as transaction:
-                    return_value = f(*args, **kwargs)
-                    transaction.set_status(SPANSTATUS.OK)
-                    return return_value
+                try:
+                    with sentry_sdk.start_span(
+                        op=OP.QUEUE_TASK_CELERY,
+                        name=task.name,
+                        source=TRANSACTION_SOURCE_TASK,
+                        origin=CeleryIntegration.origin,
+                        custom_sampling_context={
+                            "celery_job": {
+                                "task": task.name,
+                                # for some reason, args[1] is a list if non-empty but a
+                                # tuple if empty
+                                "args": list(args[1]),
+                                "kwargs": args[2],
+                            }
+                        },
+                    ) as span:
+                        return_value = f(*args, **kwargs)
+
+                except Exception:
+                    reraise(*sys.exc_info())
+                else:
+                    span.set_status(SPANSTATUS.OK)
+
+                return return_value
 
     return _inner  # type: ignore
 
@@ -368,6 +382,7 @@ def _inner(*args, **kwargs):
                 op=OP.QUEUE_PROCESS,
                 name=task.name,
                 origin=CeleryIntegration.origin,
+                only_if_parent=True,
             ) as span:
                 _set_messaging_destination_name(task, span)
 
@@ -398,12 +413,17 @@ def _inner(*args, **kwargs):
                         task.app.connection().transport.driver_type,
                     )
 
-                return f(*args, **kwargs)
+                result = f(*args, **kwargs)
+
         except Exception:
             exc_info = sys.exc_info()
             with capture_internal_exceptions():
                 _capture_exception(task, exc_info)
             reraise(*exc_info)
+        else:
+            span.set_status(SPANSTATUS.OK)
+
+        return result
 
     return _inner  # type: ignore
 
@@ -493,28 +513,36 @@ def sentry_publish(self, *args, **kwargs):
         routing_key = kwargs.get("routing_key")
         exchange = kwargs.get("exchange")
 
-        with sentry_sdk.start_span(
-            op=OP.QUEUE_PUBLISH,
-            name=task_name,
-            origin=CeleryIntegration.origin,
-            only_if_parent=True,
-        ) as span:
-            if task_id is not None:
-                span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, task_id)
-
-            if exchange == "" and routing_key is not None:
-                # Empty exchange indicates the default exchange, meaning messages are
-                # routed to the queue with the same name as the routing key.
-                span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
+        try:
+            with sentry_sdk.start_span(
+                op=OP.QUEUE_PUBLISH,
+                name=task_name,
+                origin=CeleryIntegration.origin,
+                only_if_parent=True,
+            ) as span:
+                if task_id is not None:
+                    span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, task_id)
 
-            if retries is not None:
-                span.set_data(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries)
+                if exchange == "" and routing_key is not None:
+                    # Empty exchange indicates the default exchange, meaning messages are
+                    # routed to the queue with the same name as the routing key.
+                    span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
 
-            with capture_internal_exceptions():
-                span.set_data(
-                    SPANDATA.MESSAGING_SYSTEM, self.connection.transport.driver_type
-                )
+                if retries is not None:
+                    span.set_data(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries)
+
+                with capture_internal_exceptions():
+                    span.set_data(
+                        SPANDATA.MESSAGING_SYSTEM, self.connection.transport.driver_type
+                    )
+
+                return_value = original_publish(self, *args, **kwargs)
+
+        except Exception:
+            reraise(*sys.exc_info())
+        else:
+            span.set_status(SPANSTATUS.OK)
 
-            return original_publish(self, *args, **kwargs)
+        return return_value
 
     Producer.publish = sentry_publish

tests/integrations/celery/test_celery.py

@@ -195,12 +195,12 @@ def dummy_task(x, y):
 
     events = capture_events()
 
-    with sentry_sdk.start_span(name="submission") as transaction:
+    with sentry_sdk.start_span(name="submission") as span:
         celery_invocation(dummy_task, 1, 0 if task_fails else 1)
 
     if task_fails:
         error_event = events.pop(0)
-        assert error_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+        assert error_event["contexts"]["trace"]["trace_id"] == span.trace_id
         assert error_event["exception"]["values"][0]["type"] == "ZeroDivisionError"
 
     execution_event, submission_event = events
@@ -211,8 +211,8 @@ def dummy_task(x, y):
     assert submission_event["transaction_info"] == {"source": "custom"}
 
     assert execution_event["type"] == submission_event["type"] == "transaction"
-    assert execution_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
-    assert submission_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+    assert execution_event["contexts"]["trace"]["trace_id"] == span.trace_id
+    assert submission_event["contexts"]["trace"]["trace_id"] == span.trace_id
 
     if task_fails:
         assert execution_event["contexts"]["trace"]["status"] == "internal_error"
@@ -223,7 +223,7 @@ def dummy_task(x, y):
     assert (
         execution_event["spans"][0].items()
         >= {
-            "trace_id": str(transaction.trace_id),
+            "trace_id": str(span.trace_id),
             "same_process_as_parent": True,
             "op": "queue.process",
             "description": "dummy_task",
@@ -241,7 +241,7 @@ def dummy_task(x, y):
             "span_id": submission_event["spans"][0]["span_id"],
             "start_timestamp": submission_event["spans"][0]["start_timestamp"],
             "timestamp": submission_event["spans"][0]["timestamp"],
-            "trace_id": str(transaction.trace_id),
+            "trace_id": str(span.trace_id),
         }
     ]