Fix two crashes in the aiokafka threaded producer

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,

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