diff --git a/faust/transport/drivers/aiokafka.py b/faust/transport/drivers/aiokafka.py index 8a17c1c7b..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, @@ -335,20 +336,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,20 +464,12 @@ 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 - ) - fut.set_result(ret) - return fut + # ``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) + channel = cast(ChannelT, fut.message.channel) + return await channel._finalize_message(fut, ret) else: fut2 = cast( asyncio.Future, @@ -492,10 +483,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/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: ... diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index 0f7173f4f..c05ba64be 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -1,3 +1,5 @@ +import asyncio +import inspect import random import string from contextlib import contextmanager @@ -9,6 +11,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 @@ -1897,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( @@ -1941,7 +1953,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, @@ -1957,6 +1971,95 @@ 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 = AsyncMock(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_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, + 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() + # 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( + 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()