diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index e31a3584..07363d6b 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -3,6 +3,7 @@ import contextlib import logging import multiprocessing +import os import queue import signal import threading @@ -21,6 +22,7 @@ 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.types import BrokerValue from sentry_protos.taskbroker.v1.taskbroker_pb2 import ( TASK_ACTIVATION_STATUS_COMPLETE, @@ -544,7 +546,10 @@ def check_task_future_completion( clear_current_task() processed_task_count += 1 - task_produced_futures = TaskProducer.collect_futures() + task_produced_futures = ( + TaskProducer.collect_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 if next_state != TASK_ACTIVATION_STATUS_COMPLETE: @@ -849,6 +854,9 @@ def _task_execution_complete( # Tell the parent that this child has warmed up and is ready to consume tasks messages.put_nowait(ChildMessage(child_id, "running")) + # Tell FutureTrackingProducer to track producer futures in this process + os.environ["ARROYO_TRACK_PRODUCER_FUTURES"] = "True" + # Run the worker loop run_worker( child_tasks, diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index f53cd5a7..560fbad4 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -19,6 +19,8 @@ import pytest import zstandard as zstd from arroyo.backends.kafka import KafkaPayload +from arroyo.backends.kafka.producer import FutureTrackingProducer +from arroyo.backends.kafka.producer import _pending_futures as _arroyo_pending_futures from arroyo.types import BrokerValue, Partition, Topic from redis import StrictRedis @@ -1957,16 +1959,32 @@ def test_child_process_silenced_exception_does_not_log_task_failed( assert failed_calls == [] -# Tests for TaskProducer future tracking, storage, and drain-on-shutdown behavior -# in child_process. These tests patch TaskProducer.collect_futures so we can inject +# Tests for producer future tracking, storage, and drain-on-shutdown behavior +# in child_process. These tests patch .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() @pytest.fixture @@ -2004,8 +2022,11 @@ def _producing_task(task_id: str = "task-with-futures") -> InflightTaskActivatio ) +@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_tracks_producer_futures( - clear_pending_futures: None, restore_signal_handlers: None + producer_cls: type, + clear_pending_futures: None, + restore_signal_handlers: None, ) -> None: task = _producing_task() todo: queue.Queue[InflightTaskActivation] = queue.Queue() @@ -2017,7 +2038,7 @@ def test_child_process_tracks_producer_futures( todo.put(task) with mock.patch.object( - TaskProducer, "collect_futures", return_value={"test.producer": {done_future}} + producer_cls, "collect_futures", return_value={"test.producer": {done_future}} ) as collect_mock: child_process( "examples.app:app", @@ -2039,8 +2060,11 @@ 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( - clear_pending_futures: None, restore_signal_handlers: None + producer_cls: type, + clear_pending_futures: None, + restore_signal_handlers: None, ) -> None: task = _producing_task() todo: queue.Queue[InflightTaskActivation] = queue.Queue() @@ -2066,7 +2090,7 @@ def observe_and_resolve() -> None: observer.start() try: with mock.patch.object( - TaskProducer, "collect_futures", return_value={"test.producer": {pending_future}} + producer_cls, "collect_futures", return_value={"test.producer": {pending_future}} ): child_process( "examples.app:app", @@ -2091,8 +2115,11 @@ 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( - clear_pending_futures: None, restore_signal_handlers: None + producer_cls: type, + clear_pending_futures: None, + restore_signal_handlers: None, ) -> None: task = _producing_task() todo: queue.Queue[InflightTaskActivation] = queue.Queue() @@ -2122,7 +2149,7 @@ def observe_and_resolve() -> None: observer.start() try: with mock.patch.object( - TaskProducer, "collect_futures", return_value={"test.producer": {pending_future}} + producer_cls, "collect_futures", return_value={"test.producer": {pending_future}} ): child_process( "examples.app:app", @@ -2150,8 +2177,11 @@ def observe_and_resolve() -> None: assert processed.empty() +@pytest.mark.parametrize("producer_cls", _PRODUCER_CLASSES) def test_child_process_drains_pending_futures_on_sigterm( - clear_pending_futures: None, restore_signal_handlers: None + producer_cls: type, + clear_pending_futures: None, + restore_signal_handlers: None, ) -> None: task = _producing_task() todo: queue.Queue[InflightTaskActivation] = queue.Queue() @@ -2173,7 +2203,7 @@ def deliver_sigterm() -> None: sigterm_thread.start() try: with mock.patch.object( - TaskProducer, "collect_futures", return_value={"test.producer": {pending_future}} + producer_cls, "collect_futures", return_value={"test.producer": {pending_future}} ): child_process( "examples.app:app", @@ -2195,8 +2225,11 @@ 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( - clear_pending_futures: None, restore_signal_handlers: None + producer_cls: type, + clear_pending_futures: None, + restore_signal_handlers: None, ) -> None: retriable_task = InflightTaskActivation( host="localhost:50051", @@ -2223,7 +2256,7 @@ def test_child_process_retries_on_failed_future( todo.put(retriable_task) with mock.patch.object( - TaskProducer, "collect_futures", return_value={"test.producer": {failed_future}} + producer_cls, "collect_futures", return_value={"test.producer": {failed_future}} ): child_process( "examples.app:app", @@ -2242,13 +2275,16 @@ 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( - clear_pending_futures: None, restore_signal_handlers: None + 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_futures["test.producer"].append(leftover_future) - assert len(_pending_futures) == 1 + pending_registry["test.producer"].append(leftover_future) + assert len(pending_registry) == 1 todo: queue.Queue[InflightTaskActivation] = queue.Queue() processed: queue.Queue[ProcessingResult] = queue.Queue() @@ -2274,7 +2310,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_futures) == 0 + assert len(pending_registry) == 0 def test_child_process_uses_configured_future_checking_frequency( diff --git a/uv.lock b/uv.lock index 7c49dcce..d63390b4 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.11" resolution-markers = [ "sys_platform == 'darwin' or sys_platform == 'linux'", @@ -645,13 +645,13 @@ wheels = [ [[package]] name = "sentry-arroyo" -version = "2.38.7" +version = "2.41.0" 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.38.7-py3-none-any.whl", hash = "sha256:088f8620e1fa6af950d588e4ddae259b849fb799af9a78dd5b9912ccedd19a4a" }, + { url = "https://pypi.devinfra.sentry.io/wheels/sentry_arroyo-2.41.0-py3-none-any.whl", hash = "sha256:08c0efb1a02a97ba9364f07b53a520ed204d7da015c4ea01f9bc8f9f81d6373b" }, ] [[package]] @@ -830,7 +830,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.38.7" }, + { name = "sentry-arroyo", specifier = ">=2.41.0" }, { name = "sentry-protos", specifier = ">=0.26.1" }, { name = "sentry-sdk", extras = ["http2"], specifier = ">=2.43.0" }, { name = "setuptools", marker = "extra == 'examples'", specifier = ">=80.0" },