fix(taskworker) Make task routing overrides apply at runtime (#87538)

Improve how quickly task production responds to changes in topic
routing. Instead of routing configuration only being considered during
application startup, we apply routing to each task that is produced.

Refs getsentry/taskbroker#264

Author: markstory

src/sentry/taskworker/registry.py

@@ -33,25 +33,18 @@ class TaskNamespace:
     def __init__(
         self,
         name: str,
-        topic: Topic,
+        router: TaskRouter,
         retry: Retry | None,
         expires: int | datetime.timedelta | None = None,
         processing_deadline_duration: int = DEFAULT_PROCESSING_DEADLINE,
     ):
         self.name = name
-        self.topic = topic
+        self.router = router
         self.default_retry = retry
         self.default_expires = expires  # seconds
         self.default_processing_deadline_duration = processing_deadline_duration  # seconds
         self._registered_tasks: dict[str, Task[Any, Any]] = {}
-        self._producer: SingletonProducer = SingletonProducer(
-            self._basic_producer, max_futures=1000
-        )
-
-    def _basic_producer(self) -> KafkaProducer:
-        cluster_name = get_topic_definition(self.topic)["cluster"]
-        producer_config = get_kafka_producer_cluster_options(cluster_name)
-        return KafkaProducer(producer_config)
+        self._producers: dict[Topic, SingletonProducer] = {}
 
     def get(self, name: str) -> Task[Any, Any]:
         """
@@ -69,6 +62,10 @@ def contains(self, name: str) -> bool:
         """
         return name in self._registered_tasks
 
+    @property
+    def topic(self) -> Topic:
+        return self.router.route_namespace(self.name)
+
     def register(
         self,
         *,
@@ -131,8 +128,9 @@ def wrapped(func: Callable[P, R]) -> Task[P, R]:
     def send_task(self, activation: TaskActivation, wait_for_delivery: bool = False) -> None:
         metrics.incr("taskworker.registry.send_task", tags={"namespace": activation.namespace})
 
-        produce_future = self._producer.produce(
-            ArroyoTopic(name=self.topic.value),
+        topic = self.router.route_namespace(self.name)
+        produce_future = self._producer(topic).produce(
+            ArroyoTopic(name=topic.value),
             KafkaPayload(key=None, value=activation.SerializeToString(), headers=[]),
         )
         if wait_for_delivery:
@@ -141,6 +139,17 @@ def send_task(self, activation: TaskActivation, wait_for_delivery: bool = False)
             except Exception:
                 logger.exception("Failed to wait for delivery")
 
+    def _producer(self, topic: Topic) -> SingletonProducer:
+        if topic not in self._producers:
+
+            def factory() -> KafkaProducer:
+                cluster_name = get_topic_definition(topic)["cluster"]
+                producer_config = get_kafka_producer_cluster_options(cluster_name)
+                return KafkaProducer(producer_config)
+
+            self._producers[topic] = SingletonProducer(factory, max_futures=1000)
+        return self._producers[topic]
+
 
 class TaskRegistry:
     """
@@ -191,10 +200,9 @@ def create_namespace(
 
         Namespaces can define default behavior for tasks defined within a namespace.
         """
-        topic = self._router.route_namespace(name)
         namespace = TaskNamespace(
             name=name,
-            topic=topic,
+            router=self._router,
             retry=retry,
             expires=expires,
             processing_deadline_duration=processing_deadline_duration,

src/sentry/taskworker/router.py

@@ -11,7 +11,7 @@ def route_namespace(self, name: str) -> Topic: ...
 
 
 class DefaultRouter:
-    """Simple router used for self-hosted and local development"""
+    """Router that uses django settings and options to select topics at runtime"""
 
     def route_namespace(self, name: str) -> Topic:
         overrides = options.get("taskworker.route.overrides")

src/sentry/taskworker/worker.py

@@ -467,6 +467,7 @@ def fetch_task(self) -> TaskActivation | None:
             return None
 
         if not activation:
+            metrics.incr("taskworker.worker.fetch_task.not_found")
             logger.debug("taskworker.fetch_task.not_found")
 
             self._gettask_backoff_seconds = min(self._gettask_backoff_seconds + 1, 10)

tests/sentry/taskworker/test_registry.py

@@ -1,5 +1,5 @@
 from concurrent.futures import Future
-from unittest.mock import patch
+from unittest.mock import Mock
 
 import pytest
 from django.test.utils import override_settings
@@ -11,13 +11,14 @@
 from sentry.conf.types.kafka_definition import Topic
 from sentry.taskworker.registry import TaskNamespace, TaskRegistry
 from sentry.taskworker.retry import LastAction, Retry
+from sentry.taskworker.router import DefaultRouter
 from sentry.taskworker.task import Task
 
 
 def test_namespace_register_task() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=None,
     )
 
@@ -37,7 +38,7 @@ def simple_task():
 def test_namespace_register_inherits_default_retry() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=Retry(times=5, on=(RuntimeError,)),
     )
 
@@ -65,7 +66,7 @@ def retry_none_param() -> None:
 def test_register_inherits_default_expires_processing_deadline() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=None,
         expires=10 * 60,
         processing_deadline_duration=5,
@@ -93,7 +94,7 @@ def with_expires() -> None:
 def test_namespace_get_unknown() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=None,
     )
 
@@ -102,10 +103,11 @@ def test_namespace_get_unknown() -> None:
     assert "No task registered" in str(err)
 
 
+@pytest.mark.django_db
 def test_namespace_send_task_no_retry() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=None,
     )
 
@@ -118,21 +120,24 @@ def simple_task() -> None:
     assert activation.retry_state.max_attempts == 1
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DISCARD
 
-    with patch.object(namespace, "_producer") as mock_producer:
-        namespace.send_task(activation)
-        assert mock_producer.produce.call_count == 1
+    mock_producer = Mock()
+    namespace._producers[Topic.TASK_WORKER] = mock_producer
+
+    namespace.send_task(activation)
+    assert mock_producer.produce.call_count == 1
 
-        mock_call = mock_producer.produce.call_args
-        assert mock_call[0][0].name == "task-worker"
+    mock_call = mock_producer.produce.call_args
+    assert mock_call[0][0].name == "task-worker"
 
-        proto_message = mock_call[0][1].value
-        assert proto_message == activation.SerializeToString()
+    proto_message = mock_call[0][1].value
+    assert proto_message == activation.SerializeToString()
 
 
+@pytest.mark.django_db
 def test_namespace_send_task_with_retry() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=None,
     )
 
@@ -147,19 +152,22 @@ def simple_task() -> None:
     assert activation.retry_state.max_attempts == 3
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DEADLETTER
 
-    with patch.object(namespace, "_producer") as mock_producer:
-        namespace.send_task(activation)
-        assert mock_producer.produce.call_count == 1
+    mock_producer = Mock()
+    namespace._producers[Topic.TASK_WORKER] = mock_producer
+
+    namespace.send_task(activation)
+    assert mock_producer.produce.call_count == 1
 
-        mock_call = mock_producer.produce.call_args
-        proto_message = mock_call[0][1].value
-        assert proto_message == activation.SerializeToString()
+    mock_call = mock_producer.produce.call_args
+    proto_message = mock_call[0][1].value
+    assert proto_message == activation.SerializeToString()
 
 
+@pytest.mark.django_db
 def test_namespace_with_retry_send_task() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=Retry(times=3),
     )
 
@@ -172,21 +180,24 @@ def simple_task() -> None:
     assert activation.retry_state.max_attempts == 3
     assert activation.retry_state.on_attempts_exceeded == ON_ATTEMPTS_EXCEEDED_DEADLETTER
 
-    with patch.object(namespace, "_producer") as mock_producer:
-        namespace.send_task(activation)
-        assert mock_producer.produce.call_count == 1
+    mock_producer = Mock()
+    namespace._producers[Topic.TASK_WORKER] = mock_producer
+
+    namespace.send_task(activation)
+    assert mock_producer.produce.call_count == 1
 
-        mock_call = mock_producer.produce.call_args
-        assert mock_call[0][0].name == "task-worker"
+    mock_call = mock_producer.produce.call_args
+    assert mock_call[0][0].name == "task-worker"
 
-        proto_message = mock_call[0][1].value
-        assert proto_message == activation.SerializeToString()
+    proto_message = mock_call[0][1].value
+    assert proto_message == activation.SerializeToString()
 
 
+@pytest.mark.django_db
 def test_namespace_with_wait_for_delivery_send_task() -> None:
     namespace = TaskNamespace(
         name="tests",
-        topic=Topic.TASK_WORKER,
+        router=DefaultRouter(),
         retry=Retry(times=3),
     )
 
@@ -196,18 +207,20 @@ def simple_task() -> None:
 
     activation = simple_task.create_activation()
 
-    with patch.object(namespace, "_producer") as mock_producer:
-        ret_value: Future[None] = Future()
-        ret_value.set_result(None)
-        mock_producer.produce.return_value = ret_value
-        namespace.send_task(activation, wait_for_delivery=True)
-        assert mock_producer.produce.call_count == 1
+    mock_producer = Mock()
+    namespace._producers[Topic.TASK_WORKER] = mock_producer
+
+    ret_value: Future[None] = Future()
+    ret_value.set_result(None)
+    mock_producer.produce.return_value = ret_value
+    namespace.send_task(activation, wait_for_delivery=True)
+    assert mock_producer.produce.call_count == 1
 
-        mock_call = mock_producer.produce.call_args
-        assert mock_call[0][0].name == "task-worker"
+    mock_call = mock_producer.produce.call_args
+    assert mock_call[0][0].name == "task-worker"
 
-        proto_message = mock_call[0][1].value
-        assert proto_message == activation.SerializeToString()
+    proto_message = mock_call[0][1].value
+    assert proto_message == activation.SerializeToString()
 
 
 @pytest.mark.django_db
@@ -217,7 +230,7 @@ def test_registry_get() -> None:
 
     assert isinstance(ns, TaskNamespace)
     assert ns.name == "tests"
-    assert ns.topic
+    assert ns.router
     assert ns == registry.get("tests")
 
     with pytest.raises(KeyError):
@@ -284,4 +297,6 @@ def test_registry_create_namespace_route_setting() -> None:
         assert profiling.topic == Topic.PROFILES
 
         with pytest.raises(ValueError):
-            registry.create_namespace(name="lol")
+            ns = registry.create_namespace(name="lol")
+            # Should raise as the name is routed to an invalid topic
+            ns.topic

tests/sentry/taskworker/test_task.py

@@ -7,9 +7,9 @@
     ON_ATTEMPTS_EXCEEDED_DISCARD,
 )
 
-from sentry.conf.types.kafka_definition import Topic
 from sentry.taskworker.registry import TaskNamespace
 from sentry.taskworker.retry import LastAction, Retry, RetryError
+from sentry.taskworker.router import DefaultRouter
 from sentry.taskworker.task import Task
 from sentry.testutils.helpers.task_runner import TaskRunner
 from sentry.utils import json
@@ -21,7 +21,7 @@ def do_things() -> None:
 
 @pytest.fixture
 def task_namespace() -> TaskNamespace:
-    return TaskNamespace(name="tests", topic=Topic.TASK_WORKER, retry=None)
+    return TaskNamespace(name="tests", router=DefaultRouter(), retry=None)
 
 
 def test_define_task_defaults(task_namespace: TaskNamespace) -> None: