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()