From 02a772c2397cb895c7c00b6d692407a0935e3e74 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 6 Aug 2026 12:21:25 +0000 Subject: [PATCH 1/7] Fix two crashes in the aiokafka threaded producer Both were found by the type checker in #758 and left marked `XXX` there because fixing them changes runtime behaviour. `ThreadedProducer._shutdown_thread` was a plain `def` overriding `mode.threads.ServiceThread._shutdown_thread`, which is `async def` and is awaited by `_serve()` in a `finally:`. Every shutdown of the producer thread therefore evaluated `await None` and raised TypeError. The thread only recovered because `_start_thread` catches that exception and calls `set_shutdown()` before re-raising -- so mode's teardown (`on_thread_stop`, stopping children, futures and exit stacks) never ran, and the thread died with a traceback instead of stopping cleanly. The override also scheduled `on_thread_stop()` with `asyncio.run_coroutine_threadsafe` onto `self.thread_loop` -- the loop that was about to stop, and the loop already running `_serve()`. Because the TypeError tore down `run_until_complete` immediately, that coroutine never got a chance to run, so the producer was never flushed or stopped on this path. Make it `async def` and await `super()._shutdown_thread()`, which runs `on_thread_stop()` and the rest of mode's teardown in order. The once-only guard is kept; when shutdown has already been initiated the shutdown event is still set, matching what the old TypeError path ended up doing via `_start_thread`. `ThreadedProducer.publish_message(wait=True)` called `fut.message.channel._on_published(message=..., state=..., producer=...)`. `Topic._on_published` takes the send future as a required *positional* `fut` and reads the result off it, so the call raised `TypeError: Topic._on_published() missing 1 required positional argument`. The waiting branch has no such future -- `send_and_wait` has already resolved -- so complete the message directly instead: report the sensor, set the result, and invoke the callback, which is what `Topic.publish_message(wait=True)` does via `_finalize_message`. The non-waiting branch keeps using `_on_published` as a done-callback, where `add_done_callback` supplies `fut`. `test_publish_message_with_wait` did not catch this because its channel is a bare `Mock`, which accepts any call; the new test uses a real topic and fails with the TypeError above against the previous code. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_012N2yysiNbzzVcvYhtrpsM7 --- faust/transport/drivers/aiokafka.py | 54 ++++++------ tests/unit/transport/drivers/test_aiokafka.py | 88 +++++++++++++++++++ 2 files changed, 113 insertions(+), 29 deletions(-) diff --git a/faust/transport/drivers/aiokafka.py b/faust/transport/drivers/aiokafka.py index 8a17c1c7b..2064d7535 100644 --- a/faust/transport/drivers/aiokafka.py +++ b/faust/transport/drivers/aiokafka.py @@ -335,20 +335,18 @@ def __init__( self._default_producer = default_producer self.app = app - # XXX broken: this synchronous method overrides the coroutine - # ``mode.threads.ServiceThread._shutdown_thread``, breaking mode's - # contract. ``ServiceThread._serve()`` ends with - # ``finally: await self._shutdown_thread()``, so ``await None`` raises - # TypeError on every shutdown of this thread. The base implementation - # (on_thread_stop, stopping children/futures/exit stacks, set_shutdown) - # therefore never runs; the shutdown event only gets set because - # ``_start_thread`` catches that TypeError and calls ``set_shutdown()`` - # before re-raising it. Not fixed here: making it ``async`` changes - # runtime behaviour, which is out of scope for this annotation pass. - def _shutdown_thread(self) -> None: # type: ignore[override] - # Ensure that the shutdown process is initiated only once + async def _shutdown_thread(self) -> None: + # Ensure that the shutdown process is initiated only once. + # + # This has to stay a coroutine: ``ServiceThread._serve()`` ends with + # ``finally: await self._shutdown_thread()``, so a synchronous + # override makes that ``await None`` and raises TypeError. if not self._shutdown_initiated: - asyncio.run_coroutine_threadsafe(self.on_thread_stop(), self.thread_loop) + await super()._shutdown_thread() + else: + # ``on_thread_stop`` has already run, so skip mode's teardown -- + # but still set the shutdown event, or ``stop()`` waits forever. + self.set_shutdown() async def flush(self) -> None: """Wait for producer to finish transmitting all buffered messages.""" @@ -465,19 +463,16 @@ async def publish_message( timestamp_ms=timestamp_ms, headers=headers, ) - # XXX broken: ``_on_published`` is not on the ChannelT interface - # (it is implemented by faust.topics.Topic), and worse, the call - # below is missing an argument. Topic._on_published - # (faust/topics.py:463) is - # ``_on_published(self, fut, message, producer, state)`` -- ``fut`` - # is a required positional parameter holding the send future, and - # nothing is passed for it here, so this raises TypeError at - # runtime and ``publish_message(..., wait=True)`` can never - # succeed. Behaviour left untouched in this annotation-only pass. - fut.message.channel._on_published( # type: ignore[attr-defined] - message=fut, state=state, producer=producer - ) + # ``_on_published`` is the done-callback for the non-waiting + # branch: it takes the send future positionally and reads the + # result off it. There is no such future here -- ``send_and_wait`` + # has already resolved to ``ret`` -- so complete the message + # directly, exactly as ``Topic.publish_message(wait=True)`` does + # via ``_finalize_message``. + self.app.sensors.on_send_completed(producer, state, ret) fut.set_result(ret) + if fut.message.callback: + fut.message.callback(fut) return fut else: fut2 = cast( @@ -492,10 +487,11 @@ async def publish_message( ), ) callback = partial( - # ``_on_published`` is not on the ChannelT interface; see the - # note on the ``wait`` branch above. This branch does supply - # the required positional ``fut``: add_done_callback passes - # the completed future as the first positional argument. + # ``_on_published`` is implemented by faust.topics.Topic but + # is not declared on the ChannelT interface, hence the ignore. + # Its required positional ``fut`` is supplied here by + # add_done_callback, which passes the completed send future as + # the first positional argument. fut.message.channel._on_published, # type: ignore[attr-defined] message=fut, state=state, diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index 0f7173f4f..f8b897b23 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -1,3 +1,4 @@ +import inspect import random import string from contextlib import contextmanager @@ -9,6 +10,7 @@ import pytest from aiokafka.errors import CommitFailedError, IllegalStateError, KafkaError from aiokafka.structs import OffsetAndMetadata, TopicPartition +from mode.threads import ServiceThread from mode.utils import text from mode.utils.futures import done_future from mode.utils.times import humanize_seconds_ago @@ -1957,6 +1959,92 @@ async def test_publish_message_with_wait( finally: await threaded_producer.stop() + @pytest.mark.asyncio + async def test_publish_message_with_wait__completes_the_message( + self, + *, + threaded_producer: ThreadedProducer, + mocked_producer: Mock, + app, + loop, + ): + # Regression: the wait=True branch used to call + # ``fut.message.channel._on_published(message=..., state=..., + # producer=...)``. ``Topic._on_published`` takes the send future as a + # required *positional* ``fut``, so that call raised TypeError for any + # real channel and ``publish_message(wait=True)`` could never succeed. + # ``test_publish_message_with_wait`` above does not catch it because + # its channel is a bare Mock, which accepts any call. + record_metadata = Mock(name="RecordMetadata") + mocked_producer.send_and_wait = AsyncMock(return_value=record_metadata) + threaded_producer.app.sensors = Mock(name="sensors") + callback = Mock(name="callback") + await threaded_producer.start() + try: + fut = await threaded_producer.publish_message( + wait=True, + fut_other=FutureMessage( + PendingMessage( + channel=app.topic("test-publish-wait"), + key=b"k", + value=b"v", + partition=None, + timestamp=None, + headers=None, + key_serializer=None, + value_serializer=None, + callback=callback, + ) + ), + ) + assert fut.result() is record_metadata + callback.assert_called_once_with(fut) + threaded_producer.app.sensors.on_send_completed.assert_called_once_with( + mocked_producer, + threaded_producer.app.sensors.on_send_initiated.return_value, + record_metadata, + ) + finally: + await threaded_producer.stop() + + def test_shutdown_thread_is_a_coroutine(self): + # Regression: this was a plain ``def`` overriding mode's + # ``async def ServiceThread._shutdown_thread``. ``_serve()`` ends with + # ``finally: await self._shutdown_thread()``, so a sync override makes + # that ``await None`` -- TypeError on every shutdown of the thread. + assert inspect.iscoroutinefunction(ThreadedProducer._shutdown_thread) + + @pytest.mark.asyncio + async def test_shutdown_thread__runs_mode_teardown( + self, + *, + threaded_producer: ThreadedProducer, + ): + # The old override scheduled on_thread_stop() with + # run_coroutine_threadsafe on the very loop that was about to stop, so + # mode's teardown never ran. Awaiting the base is what makes it run. + threaded_producer._shutdown_initiated = False + with patch.object(ServiceThread, "_shutdown_thread", AsyncMock()) as base: + await threaded_producer._shutdown_thread() + base.assert_called_once_with() + + @pytest.mark.asyncio + async def test_shutdown_thread__already_initiated_still_sets_shutdown( + self, + *, + threaded_producer: ThreadedProducer, + ): + threaded_producer._shutdown_initiated = True + with ( + patch.object(ServiceThread, "_shutdown_thread", AsyncMock()) as base, + patch.object(threaded_producer, "set_shutdown") as set_shutdown, + ): + await threaded_producer._shutdown_thread() + # on_thread_stop() must not run a second time, but the shutdown event + # still has to be set or stop() waits forever. + base.assert_not_called() + set_shutdown.assert_called_once_with() + class TestTransport: @pytest.fixture() From 74672648eb8c295ff9d3383431c13ca944e7105a Mon Sep 17 00:00:00 2001 From: William Barnhart Date: Wed, 12 Aug 2026 06:58:24 -0400 Subject: [PATCH 2/7] Declare channel message finalization hook --- faust/types/channels.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/faust/types/channels.py b/faust/types/channels.py index 60a2f8ad9..e7c62e63e 100644 --- a/faust/types/channels.py +++ b/faust/types/channels.py @@ -125,6 +125,11 @@ async def publish_message( self, fut: FutureMessage, wait: bool = True ) -> Awaitable[RecordMetadata]: ... + @abc.abstractmethod + async def _finalize_message( + self, fut: FutureMessage, result: RecordMetadata + ) -> FutureMessage: ... + @stampede @abc.abstractmethod async def maybe_declare(self) -> None: ... From 5c33e7fa161a921f89f0d94e2393d03bd5d8461c Mon Sep 17 00:00:00 2001 From: William Barnhart Date: Wed, 12 Aug 2026 06:58:32 -0400 Subject: [PATCH 3/7] Finalize threaded messages through Channel --- faust/transport/drivers/aiokafka.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/faust/transport/drivers/aiokafka.py b/faust/transport/drivers/aiokafka.py index 2064d7535..0c829c00f 100644 --- a/faust/transport/drivers/aiokafka.py +++ b/faust/transport/drivers/aiokafka.py @@ -82,6 +82,7 @@ from faust.types import ( TP, AppT, + ChannelT, ConsumerMessage, FutureMessage, HeadersArg, @@ -463,17 +464,12 @@ async def publish_message( timestamp_ms=timestamp_ms, headers=headers, ) - # ``_on_published`` is the done-callback for the non-waiting - # branch: it takes the send future positionally and reads the - # result off it. There is no such future here -- ``send_and_wait`` - # has already resolved to ``ret`` -- so complete the message - # directly, exactly as ``Topic.publish_message(wait=True)`` does - # via ``_finalize_message``. + # ``send_and_wait`` has already resolved the broker-side + # operation. Let the channel own Faust-level completion, + # including callback invocation and async callback handling. self.app.sensors.on_send_completed(producer, state, ret) - fut.set_result(ret) - if fut.message.callback: - fut.message.callback(fut) - return fut + channel = cast(ChannelT, fut.message.channel) + return await channel._finalize_message(fut, ret) else: fut2 = cast( asyncio.Future, From c01446df605a5f948f856c73b88ce8c436fcdbb1 Mon Sep 17 00:00:00 2001 From: William Barnhart Date: Wed, 12 Aug 2026 06:58:37 -0400 Subject: [PATCH 4/7] Test channel-owned async callback finalization --- tests/unit/transport/drivers/test_aiokafka.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index f8b897b23..43c917460 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -1978,7 +1978,7 @@ async def test_publish_message_with_wait__completes_the_message( record_metadata = Mock(name="RecordMetadata") mocked_producer.send_and_wait = AsyncMock(return_value=record_metadata) threaded_producer.app.sensors = Mock(name="sensors") - callback = Mock(name="callback") + callback = AsyncMock(name="callback") await threaded_producer.start() try: fut = await threaded_producer.publish_message( @@ -1998,7 +1998,7 @@ async def test_publish_message_with_wait__completes_the_message( ), ) assert fut.result() is record_metadata - callback.assert_called_once_with(fut) + callback.assert_awaited_once_with(fut) threaded_producer.app.sensors.on_send_completed.assert_called_once_with( mocked_producer, threaded_producer.app.sensors.on_send_initiated.return_value, From f8fe816b87c9dd60bf128e443a1edf5d926c3560 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 00:59:21 +0000 Subject: [PATCH 5/7] Fix test_publish_message_with_wait against the new wait path The wait=True branch now awaits ``channel._finalize_message(fut, ret)``, but this test still passed a bare ``Mock()`` as the channel. A plain Mock returns a Mock, which is not awaitable, so every aiokafka test leg failed with ``TypeError: object Mock can't be used in 'await' expression``. Give the stand-in channel an ``AsyncMock`` ``_finalize_message``, matching the coroutine every real channel implements. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01A5Hzidb2taZ7ci6AeBUuJT --- tests/unit/transport/drivers/test_aiokafka.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index 43c917460..ddfc6242c 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -1943,7 +1943,9 @@ async def test_publish_message_with_wait( wait=True, fut_other=FutureMessage( PendingMessage( - channel=Mock(), + # ``_finalize_message`` is a coroutine on every real + # channel, so the stand-in has to be awaitable too. + channel=Mock(_finalize_message=AsyncMock()), key="Test", value="Test", partition=None, From 50805b7ebccd87a7a75d80f89da1fb5a309dbcb3 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 01:08:38 +0000 Subject: [PATCH 6/7] Drain mode's keepalive tick in test_on_thread_stop The 3.14t free-threaded leg failed teardown with: DirtyTest: ('Left over tasks', ... " cb=[_chain_future._call_set_state()]>") That task is mode's, not faust's: ServiceThread._wakeup_timer_in_thread ends every keepalive tick with run_coroutine_threadsafe(asyncio.sleep(0), self.parent_loop) and fires one last tick as the thread stops. The task completes on the next iteration of the parent loop -- but the autouse tasks_not_lingering fixture snapshots tasks as soon as the test coroutine returns, so if the loop never runs again it is still pending and gets reported. It surfaced here because _shutdown_thread now performs mode's real teardown instead of raising TypeError out of _serve(), which gives the keepalive room to tick before the thread goes away. Yield once after stop() so the transient task settles. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01A5Hzidb2taZ7ci6AeBUuJT --- tests/unit/transport/drivers/test_aiokafka.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index ddfc6242c..90436830b 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -1,3 +1,4 @@ +import asyncio import inspect import random import string @@ -1899,6 +1900,15 @@ async def test_on_thread_stop( mocked_producer.stop.assert_called_once() finally: await threaded_producer.stop() + # mode's thread keepalive (``ServiceThread._wakeup_timer_in_thread``) + # ends each tick with ``run_coroutine_threadsafe(asyncio.sleep(0), + # parent_loop)``, and it gets one last tick out of the way as the + # thread stops. That leaves a transient ``sleep()`` task on the + # parent loop which the autouse ``tasks_not_lingering`` fixture + # reports as "Left over tasks" if the loop never runs again -- + # it did on the 3.14t leg, where the now-real teardown gives the + # keepalive room to fire. One iteration is all the task needs. + await asyncio.sleep(0.1) @pytest.mark.asyncio async def test_publish_message( From fdbf5a88b6c1e5c11ab939c3ea32b2a8def9e4bf Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 18:17:05 +0000 Subject: [PATCH 7/7] Assert the base shutdown coroutine is awaited, not just called Review feedback: assert_called_once_with() passes for a caller that invokes the coroutine without awaiting it -- which is exactly the bug this test guards against. assert_awaited_once_with() checks the await. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01A5Hzidb2taZ7ci6AeBUuJT --- tests/unit/transport/drivers/test_aiokafka.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index 90436830b..c05ba64be 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -2038,7 +2038,10 @@ async def test_shutdown_thread__runs_mode_teardown( threaded_producer._shutdown_initiated = False with patch.object(ServiceThread, "_shutdown_thread", AsyncMock()) as base: await threaded_producer._shutdown_thread() - base.assert_called_once_with() + # Awaited, not merely called: the bug being guarded against here is a + # coroutine that never gets awaited, which assert_called_once_with + # would happily accept. + base.assert_awaited_once_with() @pytest.mark.asyncio async def test_shutdown_thread__already_initiated_still_sets_shutdown(