Fix two crashes in the aiokafka threaded producer

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,

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: ...

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