Propagate sentry-trace and baggage to huey tasks (#2792)

This PR enables passing `sentry-trace` and `baggage` headers to background tasks using the Huey task queue.

This allows easily correlating what happens inside a background task with whatever transaction (e.g. a user request in a Django application) queued the task in the first place.

Periodic tasks do not get these headers, because otherwise each execution of the periodic task would be tied to the same parent trace (the long-running worker process).

--- 

Co-authored-by: Anton Pirker <[email protected]>

Author: cnschn

sentry_sdk/integrations/huey.py

@@ -6,10 +6,15 @@
 from sentry_sdk._compat import reraise
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk import Hub
+from sentry_sdk.api import continue_trace, get_baggage, get_traceparent
 from sentry_sdk.consts import OP
 from sentry_sdk.hub import _should_send_default_pii
 from sentry_sdk.integrations import DidNotEnable, Integration
-from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_TASK
+from sentry_sdk.tracing import (
+    BAGGAGE_HEADER_NAME,
+    SENTRY_TRACE_HEADER_NAME,
+    TRANSACTION_SOURCE_TASK,
+)
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -25,7 +30,7 @@
     F = TypeVar("F", bound=Callable[..., Any])
 
 try:
-    from huey.api import Huey, Result, ResultGroup, Task
+    from huey.api import Huey, Result, ResultGroup, Task, PeriodicTask
     from huey.exceptions import CancelExecution, RetryTask, TaskLockedException
 except ImportError:
     raise DidNotEnable("Huey is not installed")
@@ -56,6 +61,14 @@ def _sentry_enqueue(self, task):
             return old_enqueue(self, task)
 
         with hub.start_span(op=OP.QUEUE_SUBMIT_HUEY, description=task.name):
+            if not isinstance(task, PeriodicTask):
+                # Attach trace propagation data to task kwargs. We do
+                # not do this for periodic tasks, as these don't
+                # really have an originating transaction.
+                task.kwargs["sentry_headers"] = {
+                    BAGGAGE_HEADER_NAME: get_baggage(),
+                    SENTRY_TRACE_HEADER_NAME: get_traceparent(),
+                }
             return old_enqueue(self, task)
 
     Huey.enqueue = _sentry_enqueue
@@ -145,12 +158,15 @@ def _sentry_execute(self, task, timestamp=None):
                 scope.clear_breadcrumbs()
                 scope.add_event_processor(_make_event_processor(task))
 
-            transaction = Transaction(
+            sentry_headers = task.kwargs.pop("sentry_headers", None)
+
+            transaction = continue_trace(
+                sentry_headers or {},
                 name=task.name,
-                status="ok",
                 op=OP.QUEUE_TASK_HUEY,
                 source=TRANSACTION_SOURCE_TASK,
             )
+            transaction.set_status("ok")
 
             if not getattr(task, "_sentry_is_patched", False):
                 task.execute = _wrap_task_execute(task.execute)

tests/integrations/huey/test_huey.py

@@ -172,3 +172,21 @@ def dummy_task():
     assert len(event["spans"])
     assert event["spans"][0]["op"] == "queue.submit.huey"
     assert event["spans"][0]["description"] == "different_task_name"
+
+
+def test_huey_propagate_trace(init_huey, capture_events):
+    huey = init_huey()
+
+    events = capture_events()
+
+    @huey.task()
+    def propagated_trace_task():
+        pass
+
+    with start_transaction() as outer_transaction:
+        execute_huey_task(huey, propagated_trace_task)
+
+    assert (
+        events[0]["transaction"] == "propagated_trace_task"
+    )  # the "inner" transaction
+    assert events[0]["contexts"]["trace"]["trace_id"] == outer_transaction.trace_id