feat(taskworker) Add support for headers to Task.apply_async() (#88689)

We use task headers in a few places and supporting it is simple.

Refs #88462

Author: markstory

src/sentry/taskworker/task.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import datetime
-from collections.abc import Callable
+from collections.abc import Callable, Collection, Mapping
 from functools import update_wrapper
 from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar
 from uuid import uuid4
@@ -81,30 +81,32 @@ def delay(self, *args: P.args, **kwargs: P.kwargs) -> None:
         The provided parameters will be JSON encoded and stored within
         a `TaskActivation` protobuf that is appended to kafka
         """
-        if settings.TASK_WORKER_ALWAYS_EAGER:
-            self._func(*args, **kwargs)
-        else:
-            # TODO(taskworker) promote parameters to headers
-            self._namespace.send_task(
-                self.create_activation(*args, **kwargs), wait_for_delivery=self.wait_for_delivery
-            )
+        self.apply_async(args=args, kwargs=kwargs)
 
-    def apply_async(self, args: Any = None, kwargs: Any = None) -> None:
+    def apply_async(
+        self, args: Any = None, kwargs: Any = None, headers: Mapping[str, Any] | None = None
+    ) -> None:
         """
         Schedule a task to run later with a set of arguments.
 
         The provided parameters will be JSON encoded and stored within
         a `TaskActivation` protobuf that is appended to kafka
-
-        Prefer using `delay()` instead of `apply_async()`.
         """
-        if args is None:
-            args = []
-        if kwargs is None:
-            kwargs = {}
-        self.delay(*args, **kwargs)
+        if settings.TASK_WORKER_ALWAYS_EAGER:
+            self._func(*args, **kwargs)
+        else:
+            # TODO(taskworker) promote parameters to headers
+            self._namespace.send_task(
+                self.create_activation(args=args, kwargs=kwargs, headers=headers),
+                wait_for_delivery=self.wait_for_delivery,
+            )
 
-    def create_activation(self, *args: P.args, **kwargs: P.kwargs) -> TaskActivation:
+    def create_activation(
+        self,
+        args: Collection[Any],
+        kwargs: Mapping[Any, Any],
+        headers: Mapping[str, Any] | None = None,
+    ) -> TaskActivation:
         received_at = Timestamp()
         received_at.FromDatetime(timezone.now())
 
@@ -116,9 +118,13 @@ def create_activation(self, *args: P.args, **kwargs: P.kwargs) -> TaskActivation
         if isinstance(expires, datetime.timedelta):
             expires = int(expires.total_seconds())
 
+        if not headers:
+            headers = {}
+
         headers = {
             "sentry-trace": sentry_sdk.get_traceparent() or "",
             "baggage": sentry_sdk.get_baggage() or "",
+            **headers,
         }
 
         return TaskActivation(

tests/sentry/taskworker/test_registry.py

@@ -81,12 +81,12 @@ def with_expires() -> None:
         raise NotImplementedError
 
     no_expires_task = namespace.get("test.no_expires")
-    activation = no_expires_task.create_activation()
+    activation = no_expires_task.create_activation([], {})
     assert activation.expires == 10 * 60
     assert activation.processing_deadline_duration == 5
 
     with_expires_task = namespace.get("test.with_expires")
-    activation = with_expires_task.create_activation()
+    activation = with_expires_task.create_activation([], {})
     assert activation.expires == 30 * 60
     assert activation.processing_deadline_duration == 10
 
@@ -115,7 +115,7 @@ def test_namespace_send_task_no_retry() -> None:
     def simple_task() -> None:
         raise NotImplementedError
 
-    activation = simple_task.create_activation()
+    activation = simple_task.create_activation([], {})
     assert activation.retry_state.attempts == 0
     assert activation.retry_state.max_attempts == 1
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DISCARD
@@ -147,7 +147,7 @@ def test_namespace_send_task_with_retry() -> None:
     def simple_task() -> None:
         raise NotImplementedError
 
-    activation = simple_task.create_activation()
+    activation = simple_task.create_activation([], {})
     assert activation.retry_state.attempts == 0
     assert activation.retry_state.max_attempts == 3
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DEADLETTER
@@ -175,7 +175,7 @@ def test_namespace_with_retry_send_task() -> None:
     def simple_task() -> None:
         raise NotImplementedError
 
-    activation = simple_task.create_activation()
+    activation = simple_task.create_activation([], {})
     assert activation.retry_state.attempts == 0
     assert activation.retry_state.max_attempts == 3
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DEADLETTER
@@ -205,7 +205,7 @@ def test_namespace_with_wait_for_delivery_send_task() -> None:
     def simple_task() -> None:
         raise NotImplementedError
 
-    activation = simple_task.create_activation()
+    activation = simple_task.create_activation([], {})
 
     mock_producer = Mock()
     namespace._producers[Topic.TASKWORKER] = mock_producer

tests/sentry/taskworker/test_task.py

@@ -133,33 +133,33 @@ def test_create_activation(task_namespace: TaskNamespace) -> None:
         at_most_once=True,
     )
     # No retries will be made as there is no retry policy on the task or namespace.
-    activation = no_retry_task.create_activation()
+    activation = no_retry_task.create_activation([], {})
     assert activation.taskname == "test.no_retry"
     assert activation.namespace == task_namespace.name
     assert activation.retry_state
     assert activation.retry_state.attempts == 0
     assert activation.retry_state.max_attempts == 1
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DISCARD
 
-    activation = retry_task.create_activation()
+    activation = retry_task.create_activation([], {})
     assert activation.taskname == "test.with_retry"
     assert activation.namespace == task_namespace.name
     assert activation.retry_state
     assert activation.retry_state.attempts == 0
     assert activation.retry_state.max_attempts == 3
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DEADLETTER
 
-    activation = timedelta_expiry_task.create_activation()
+    activation = timedelta_expiry_task.create_activation([], {})
     assert activation.taskname == "test.with_timedelta_expires"
     assert activation.expires == 300
     assert activation.processing_deadline_duration == 30
 
-    activation = int_expiry_task.create_activation()
+    activation = int_expiry_task.create_activation([], {})
     assert activation.taskname == "test.with_int_expires"
     assert activation.expires == 300
     assert activation.processing_deadline_duration == 30
 
-    activation = at_most_once_task.create_activation()
+    activation = at_most_once_task.create_activation([], {})
     assert activation.taskname == "test.at_most_once"
     assert activation.namespace == task_namespace.name
     assert activation.retry_state
@@ -174,7 +174,7 @@ def test_create_activation_parameters(task_namespace: TaskNamespace) -> None:
     def with_parameters(one: str, two: int, org_id: int) -> None:
         raise NotImplementedError
 
-    activation = with_parameters.create_activation("one", 22, org_id=99)
+    activation = with_parameters.create_activation(["one", 22], {"org_id": 99})
     params = json.loads(activation.parameters)
     assert params["args"]
     assert params["args"] == ["one", 22]
@@ -187,8 +187,24 @@ def with_parameters(one: str, two: int, org_id: int) -> None:
         raise NotImplementedError
 
     with sentry_sdk.start_transaction(op="test.task"):
-        activation = with_parameters.create_activation("one", 22, org_id=99)
+        activation = with_parameters.create_activation(["one", 22], {"org_id": 99})
 
     headers = activation.headers
     assert headers["sentry-trace"]
     assert "baggage" in headers
+
+
+def test_create_activation_headers(task_namespace: TaskNamespace) -> None:
+    @task_namespace.register(name="test.parameters")
+    def with_parameters(one: str, two: int, org_id: int) -> None:
+        raise NotImplementedError
+
+    with sentry_sdk.start_transaction(op="test.task"):
+        activation = with_parameters.create_activation(
+            ["one", 22], {"org_id": 99}, {"key": "value"}
+        )
+
+    headers = activation.headers
+    assert headers["sentry-trace"]
+    assert "baggage" in headers
+    assert headers["key"] == "value"