Skip to content

feat(workerchild): begin porting producer to FutureTrackingProducer #744

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bmckerry merged 1 commit into from
Jul 6, 2026
Merged

feat(workerchild): begin porting producer to FutureTrackingProducer #744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion clients/python/src/taskbroker_client/worker/workerchild.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import contextlib
import logging
import multiprocessing
import os
import queue
import signal
import threading
Expand All @@ -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,
Expand Down Expand Up @@ -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()
)

@cursor cursor Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicate keys drop futures

Medium Severity

Merging TaskProducer.collect_futures() and FutureTrackingProducer.collect_futures() with dict | keeps only one set of futures when both registries share a producer name. TaskProducer entries are overwritten by FutureTrackingProducer, so those Kafka futures are never awaited and the activation can complete without them finishing or failing.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 0700ac1. Configure here.

@bmckerry bmckerry Jul 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not an issue, we won't share names during the cutover


# 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:
Expand Down Expand Up @@ -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,
Expand Down
68 changes: 52 additions & 16 deletions clients/python/tests/worker/test_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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 <producer>.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
Expand Down Expand Up @@ -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()
Expand All @@ -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",
Expand All @@ -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()
Expand All @@ -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",
Expand All @@ -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()
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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()
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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()
Expand All @@ -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(
Expand Down
8 changes: 4 additions & 4 deletions uv.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading