From 6043d213612778315ccb014130fa43903728f2b7 Mon Sep 17 00:00:00 2001 From: Ben McKerry <110857332+bmckerry@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:23:44 -0400 Subject: [PATCH] feat(worker): remove TaskProducer --- clients/python/pyproject.toml | 2 +- clients/python/src/examples/tasks.py | 5 +- .../src/taskbroker_client/worker/producer.py | 105 ------------------ .../taskbroker_client/worker/workerchild.py | 8 +- clients/python/tests/worker/test_producer.py | 96 ---------------- clients/python/tests/worker/test_worker.py | 53 +++------ uv.lock | 6 +- 7 files changed, 23 insertions(+), 252 deletions(-) delete mode 100644 clients/python/src/taskbroker_client/worker/producer.py delete mode 100644 clients/python/tests/worker/test_producer.py diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index 92656511..b9e231ce 100644 --- a/clients/python/pyproject.toml +++ b/clients/python/pyproject.toml @@ -4,7 +4,7 @@ version = "0.20.12" description = "Taskbroker python client and worker runtime" readme = "README.md" dependencies = [ - "sentry-arroyo>=2.41.0", + "sentry-arroyo>=2.41.1", "sentry-sdk[http2]>=2.43.0", "sentry-protos>=0.26.1", "confluent_kafka>=2.3.0", diff --git a/clients/python/src/examples/tasks.py b/clients/python/src/examples/tasks.py index b1f90e5d..c3474660 100644 --- a/clients/python/src/examples/tasks.py +++ b/clients/python/src/examples/tasks.py @@ -9,14 +9,13 @@ from time import sleep from typing import Any -from arroyo.backends.kafka import KafkaPayload, KafkaProducer +from arroyo.backends.kafka import FutureTrackingProducer, KafkaPayload, KafkaProducer from arroyo.types import Topic from redis import StrictRedis from examples.app import app from taskbroker_client.retry import LastAction, NoRetriesRemainingError, Retry, RetryTaskError from taskbroker_client.retry import retry_task as retry_task_helper -from taskbroker_client.worker.producer import TaskProducer from taskbroker_client.worker.workerchild import ProcessingDeadlineExceeded logger = logging.getLogger(__name__) @@ -135,7 +134,7 @@ def task_that_produces( def producer_factory() -> KafkaProducer: return KafkaProducer({"bootstrap.servers": bootstrap_servers}) - producer = TaskProducer("test.producer", producer_factory) + producer = FutureTrackingProducer("test.producer", producer_factory) production_count = random.randint(1, 50) if random_count else production_count for i in range(production_count): logger.debug(f"Producing message {i} onto topic {destination_topic}...") diff --git a/clients/python/src/taskbroker_client/worker/producer.py b/clients/python/src/taskbroker_client/worker/producer.py deleted file mode 100644 index 71d6a496..00000000 --- a/clients/python/src/taskbroker_client/worker/producer.py +++ /dev/null @@ -1,105 +0,0 @@ -import atexit -from collections import defaultdict, deque -from collections.abc import Callable -from concurrent.futures import Future -from typing import Any, Sequence - -from arroyo.backends.abstract import ProducerFuture, SimpleProducerFuture -from arroyo.backends.kafka import KafkaPayload -from arroyo.types import BrokerValue, Partition, Topic - -from taskbroker_client.constants import TASK_PRODUCER_MAX_PENDING_FUTURES -from taskbroker_client.metrics import MetricsBackend, NoOpMetricsBackend -from taskbroker_client.types import CloseableProducerProtocol - -# This is global as TaskWorker needs to be able to call TaskProducer.collect_futures() -# without a reference to a task's specific instance of TaskProducer. -# Keys are the names of each `TaskProducer` instance in the current process, values are -# deques with a maxlen to prevent unbounded queue size if `collect_futures()` is never called. -_pending_futures: defaultdict[str, deque[ProducerFuture[BrokerValue[KafkaPayload]]]] = defaultdict( - lambda: deque(maxlen=TASK_PRODUCER_MAX_PENDING_FUTURES) -) - - -class TaskProducer: - """ - TaskProducer is a producer abstraction that should be used by tasks - that produce to Kafka as a side effect of their task function. - After a TaskWorker child process executes a task activation, it will collect all - producer futures tracked by TaskProducer, and will only register the task activation as - a success if all producer futures from that activation were successful. - Otherwise, the activation will be retried. - - Args: - name: Unique identifying name of this TaskProducer. Used in metric tags. - producer_factory: Callable that returns a producer object. - metrics_backend: Application metrics backend this producer should use. - Defaults to NoOpMetricsBackend. - """ - - def __init__( - self, - name: str, - producer_factory: Callable[[], CloseableProducerProtocol], - metrics_backend: MetricsBackend | None = None, - ) -> None: - self.name = name - self._producer_factory = producer_factory - self._inner_producer: CloseableProducerProtocol | None = None - self.metrics = metrics_backend if metrics_backend is not None else NoOpMetricsBackend() - - def _get(self) -> CloseableProducerProtocol: - if self._inner_producer is None: - self._inner_producer = self._producer_factory() - atexit.register(self._shutdown) - return self._inner_producer - - def track_future(self, future: ProducerFuture[BrokerValue[KafkaPayload]]) -> None: - _pending_futures[self.name].append(future) - - @staticmethod - def collect_futures() -> dict[str, set[ProducerFuture[BrokerValue[KafkaPayload]]]]: - """ - Clears the `_pending_futures` dict, and returns a copy with all values converted to sets. - """ - pending_copy = _pending_futures.copy() - _pending_futures.clear() - return {name: set(val) for name, val in pending_copy.items()} - - def produce( - self, - dest: Topic | Partition, - payload: KafkaPayload, - callbacks: Sequence[Callable[[Future[BrokerValue[KafkaPayload]]], Any]] = [], - ) -> None: - """ - Produces the given payload to the given topic. - Since TaskProducer tracks futures internally, it does not return the - producer future to the user, but the user can still add callbacks - to the future via the `callbacks` arg. - - Args: - dest: Topic (or specific partition) to produce to. - payload: KafkaPayload to produce. - callbacks: List of Callables to add to the future as done callbacks. The future itself - is the only arg passed to the callback. - """ - future = self._get().produce(dest, payload) - self.track_future(future) - if callbacks: - # Arroyo producers can return a SimpleProducerFuture, - # which does not accept callbacks. - if not isinstance(future, SimpleProducerFuture): - for c in callbacks: - future.add_done_callback(c) - else: - raise RuntimeError( - ( - "Cannot add callbacks to SimpleProducerFuture, either remove the callbacks " - "or instantiate your producer with `use_simple_futures=False`." - ) - ) - - def _shutdown(self) -> None: - if self._inner_producer is not None: - self._inner_producer.close().result() diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index d8f206e1..6b57b928 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -20,8 +20,7 @@ import sentry_sdk import zstandard as zstd from arroyo.backends.abstract import ProducerFuture -from arroyo.backends.kafka import KafkaPayload -from arroyo.backends.kafka.producer import FutureTrackingProducer +from arroyo.backends.kafka import FutureTrackingProducer, KafkaPayload from arroyo.types import BrokerValue from sentry_protos.taskbroker.v1.taskbroker_pb2 import ( TASK_ACTIVATION_STATUS_COMPLETE, @@ -39,7 +38,6 @@ from taskbroker_client.state import clear_current_task, current_task, set_current_task from taskbroker_client.task import Task from taskbroker_client.types import ContextHook, InflightTaskActivation, ProcessingResult -from taskbroker_client.worker.producer import TaskProducer logger = logging.getLogger(__name__) @@ -547,9 +545,7 @@ def check_task_future_completion( # To have Taskworker track futures, set the env var `ARROYO_TRACK_PRODUCER_FUTURES = True` # in the worker process - task_produced_futures = ( - TaskProducer.collect_futures() | FutureTrackingProducer.collect_futures() - ) + task_produced_futures = FutureTrackingProducer.collect_futures() # If the task function itself failed, we don't need to await any # producer futures since it'll be retried anyways diff --git a/clients/python/tests/worker/test_producer.py b/clients/python/tests/worker/test_producer.py deleted file mode 100644 index afb6e52c..00000000 --- a/clients/python/tests/worker/test_producer.py +++ /dev/null @@ -1,96 +0,0 @@ -from collections.abc import Iterator -from concurrent.futures import Future -from datetime import datetime -from functools import partial - -import pytest -from arroyo.backends.abstract import ProducerFuture, SimpleProducerFuture -from arroyo.backends.kafka import KafkaPayload -from arroyo.types import BrokerValue, Partition, Topic - -from taskbroker_client.worker.producer import TaskProducer, _pending_futures - - -def make_kafka_payload() -> KafkaPayload: - """Generates dummy KafkaPayload.""" - return KafkaPayload(None, b"", []) - - -def make_broker_value() -> BrokerValue[KafkaPayload]: - """Generates dummy BrokerValue[KafkaPayload].""" - return BrokerValue(make_kafka_payload(), Partition(Topic("test"), 0), 0, datetime(1999, 2, 19)) - - -class DummyProducer: - def __init__(self, use_simple_futures: bool): - self.use_simple_futures = use_simple_futures - - def produce( - self, destination: Topic | Partition, payload: KafkaPayload - ) -> ProducerFuture[BrokerValue[KafkaPayload]]: - future: ProducerFuture[BrokerValue[KafkaPayload]] - if self.use_simple_futures: - future = SimpleProducerFuture() - else: - future = Future() - future.set_result(make_broker_value()) - return future - - def close(self) -> Future[None]: - f: Future[None] = Future() - f.set_result(None) - return f - - -def get_dummy_producer(use_simple_futures: bool) -> DummyProducer: - return DummyProducer(use_simple_futures=use_simple_futures) - - -@pytest.fixture(autouse=True) -def clear_pending_futures() -> Iterator[None]: - _pending_futures.clear() - yield - _pending_futures.clear() - - -def test_producer_tracks_futures() -> None: - producer = TaskProducer("test.producer", partial(get_dummy_producer, use_simple_futures=True)) - producer.produce(Topic("test"), make_kafka_payload()) - assert len(_pending_futures) == 1 - collected = TaskProducer.collect_futures() - future = next(iter(collected["test.producer"])) - assert future.result() == make_broker_value() - assert len(_pending_futures) == 0 - - -def test_producer_executes_callbacks() -> None: - producer = TaskProducer("test.producer", partial(get_dummy_producer, use_simple_futures=False)) - received: list[Future[BrokerValue[KafkaPayload]]] = [] - - def callback(future: Future[BrokerValue[KafkaPayload]]) -> None: - received.append(future) - - producer.produce(Topic("test"), make_kafka_payload(), callbacks=[callback]) - collected = TaskProducer.collect_futures() - tracked_future = next(iter(collected["test.producer"])) - - assert len(received) == 1 - assert received[0] is tracked_future - assert received[0].done() - - -def test_producer_rejects_callbacks_for_simple_futures() -> None: - producer = TaskProducer("test.producer", partial(get_dummy_producer, use_simple_futures=True)) - - def callback(future: Future[BrokerValue[KafkaPayload]]) -> None: - pass - - with pytest.raises(RuntimeError, match="SimpleProducerFuture"): - producer.produce(Topic("test"), make_kafka_payload(), callbacks=[callback]) - - -def test_pending_futures_max_len() -> None: - producer = TaskProducer("test.producer", partial(get_dummy_producer, use_simple_futures=True)) - for _ in range(10001): - producer.produce(Topic("test"), make_kafka_payload()) - assert len(_pending_futures["test.producer"]) == 10000 diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 560fbad4..2f5d8796 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -42,7 +42,6 @@ from taskbroker_client.retry import NoRetriesRemainingError from taskbroker_client.state import current_task from taskbroker_client.types import InflightTaskActivation, ProcessingResult -from taskbroker_client.worker.producer import TaskProducer, _pending_futures from taskbroker_client.worker.worker import ( BatchPushTaskWorker, PushTaskWorker, @@ -1960,30 +1959,14 @@ def test_child_process_silenced_exception_does_not_log_task_failed( # Tests for producer future tracking, storage, and drain-on-shutdown behavior -# in child_process. These tests patch .collect_futures so we can inject +# in child_process. These tests patch FutureTrackingProducer.collect_futures so we can inject # controllable futures without needing a real Kafka broker. -# -# child_process collects futures from both the local TaskProducer and arroyo's -# FutureTrackingProducer (unioning the two registries), so the tests are -# parametrized to run identically against either producer. This will be removed -# once all clients are fully ported from TaskProducer to FutureTrackingProducer. -_PRODUCER_CLASSES = [ - pytest.param(TaskProducer, id="task_producer"), - pytest.param(FutureTrackingProducer, id="future_tracking_producer"), -] - -_PENDING_REGISTRIES = [ - pytest.param(_pending_futures, id="task_producer"), - pytest.param(_arroyo_pending_futures, id="future_tracking_producer"), -] @pytest.fixture def clear_pending_futures() -> Iterator[None]: - _pending_futures.clear() _arroyo_pending_futures.clear() yield - _pending_futures.clear() _arroyo_pending_futures.clear() @@ -2022,9 +2005,7 @@ def _producing_task(task_id: str = "task-with-futures") -> InflightTaskActivatio ) -@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_tracks_producer_futures( - producer_cls: type, clear_pending_futures: None, restore_signal_handlers: None, ) -> None: @@ -2038,7 +2019,7 @@ def test_child_process_tracks_producer_futures( todo.put(task) with mock.patch.object( - producer_cls, "collect_futures", return_value={"test.producer": {done_future}} + FutureTrackingProducer, "collect_futures", return_value={"test.producer": {done_future}} ) as collect_mock: child_process( "examples.app:app", @@ -2060,9 +2041,7 @@ def test_child_process_tracks_producer_futures( assert result.status == TASK_ACTIVATION_STATUS_COMPLETE -@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_holds_result_until_futures_done( - producer_cls: type, clear_pending_futures: None, restore_signal_handlers: None, ) -> None: @@ -2090,7 +2069,9 @@ def observe_and_resolve() -> None: observer.start() try: with mock.patch.object( - producer_cls, "collect_futures", return_value={"test.producer": {pending_future}} + FutureTrackingProducer, + "collect_futures", + return_value={"test.producer": {pending_future}}, ): child_process( "examples.app:app", @@ -2115,9 +2096,7 @@ def observe_and_resolve() -> None: assert result.status == TASK_ACTIVATION_STATUS_COMPLETE -@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_skip_awaiting_futures_places_result_immediately( - producer_cls: type, clear_pending_futures: None, restore_signal_handlers: None, ) -> None: @@ -2149,7 +2128,9 @@ def observe_and_resolve() -> None: observer.start() try: with mock.patch.object( - producer_cls, "collect_futures", return_value={"test.producer": {pending_future}} + FutureTrackingProducer, + "collect_futures", + return_value={"test.producer": {pending_future}}, ): child_process( "examples.app:app", @@ -2177,9 +2158,7 @@ def observe_and_resolve() -> None: assert processed.empty() -@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_drains_pending_futures_on_sigterm( - producer_cls: type, clear_pending_futures: None, restore_signal_handlers: None, ) -> None: @@ -2203,7 +2182,9 @@ def deliver_sigterm() -> None: sigterm_thread.start() try: with mock.patch.object( - producer_cls, "collect_futures", return_value={"test.producer": {pending_future}} + FutureTrackingProducer, + "collect_futures", + return_value={"test.producer": {pending_future}}, ): child_process( "examples.app:app", @@ -2225,9 +2206,7 @@ def deliver_sigterm() -> None: assert result.status == TASK_ACTIVATION_STATUS_COMPLETE -@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_retries_on_failed_future( - producer_cls: type, clear_pending_futures: None, restore_signal_handlers: None, ) -> None: @@ -2256,7 +2235,7 @@ def test_child_process_retries_on_failed_future( todo.put(retriable_task) with mock.patch.object( - producer_cls, "collect_futures", return_value={"test.producer": {failed_future}} + FutureTrackingProducer, "collect_futures", return_value={"test.producer": {failed_future}} ): child_process( "examples.app:app", @@ -2275,16 +2254,14 @@ def test_child_process_retries_on_failed_future( assert result.status == TASK_ACTIVATION_STATUS_RETRY -@pytest.mark.parametrize("pending_registry", _PENDING_REGISTRIES) def test_child_process_clears_pending_futures_when_task_fails( - pending_registry: Any, clear_pending_futures: None, restore_signal_handlers: None, ) -> None: leftover_future: Future[BrokerValue[KafkaPayload]] = Future() leftover_future.set_result(_make_broker_value()) - pending_registry["test.producer"].append(leftover_future) - assert len(pending_registry) == 1 + _arroyo_pending_futures["test.producer"].append(leftover_future) + assert len(_arroyo_pending_futures) == 1 todo: queue.Queue[InflightTaskActivation] = queue.Queue() processed: queue.Queue[ProcessingResult] = queue.Queue() @@ -2310,7 +2287,7 @@ def test_child_process_clears_pending_futures_when_task_fails( # The orphaned future is dropped (the activation will be retried at the # broker level if applicable) but the global registry is cleared so it # cannot bleed into the next task this child processes. - assert len(pending_registry) == 0 + assert len(_arroyo_pending_futures) == 0 def test_child_process_uses_configured_future_checking_frequency( diff --git a/uv.lock b/uv.lock index 4776d70b..f850188e 100644 --- a/uv.lock +++ b/uv.lock @@ -667,13 +667,13 @@ wheels = [ [[package]] name = "sentry-arroyo" -version = "2.41.0" +version = "2.41.1" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ { name = "confluent-kafka", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] wheels = [ - { url = "https://pypi.devinfra.sentry.io/wheels/sentry_arroyo-2.41.0-py3-none-any.whl", hash = "sha256:08c0efb1a02a97ba9364f07b53a520ed204d7da015c4ea01f9bc8f9f81d6373b" }, + { url = "https://pypi.devinfra.sentry.io/wheels/sentry_arroyo-2.41.1-py3-none-any.whl", hash = "sha256:12465dd0d388d2ad4abab5e243ca348ad22bd6dc68434d765e41683c9a565573" }, ] [[package]] @@ -854,7 +854,7 @@ requires-dist = [ { name = "protobuf", specifier = ">=5.28.3" }, { name = "redis", specifier = ">=3.4.1" }, { name = "redis-py-cluster", marker = "extra == 'cluster'", specifier = ">=2.1.0" }, - { name = "sentry-arroyo", specifier = ">=2.41.0" }, + { name = "sentry-arroyo", specifier = ">=2.41.1" }, { name = "sentry-protos", specifier = ">=0.26.1" }, { name = "sentry-sdk", extras = ["http2"], specifier = ">=2.43.0" }, { name = "setuptools", marker = "extra == 'examples'", specifier = ">=80.0" },