Keep threaded sends inside exactly-once transactions

faust/transport/producer.py

@@ -52,10 +52,11 @@ def put(self, fut: FutureMessage) -> None:
         # both swap the AttributeError raised by a half-built buffer for an
         # AssertionError and disappear under ``python -O``.
         app = cast(AppT, self.app)
-        if app.conf.producer_threaded:
+        threaded_producer = getattr(self, "threaded_producer", None)
+        if threaded_producer is not None and not app.in_transaction:
             # Likewise Optional: only set when :setting:`producer_threaded`
             # is enabled, which is the branch we are in.
-            threaded_producer = cast(ServiceThread, self.threaded_producer)
+            threaded_producer = cast(ServiceThread, threaded_producer)
             if not self.queue:
                 # ``event_queue`` is not part of the ``ServiceThread``
                 # interface: it belongs to the driver's threaded producer
@@ -132,7 +133,8 @@ def size(self) -> int:
         # after construction; ``cast`` keeps the original AttributeError
         # behaviour instead of an assert that ``python -O`` would remove.
         app = cast(AppT, self.app)
-        if app.conf.producer_threaded:
+        threaded_producer = getattr(self, "threaded_producer", None)
+        if threaded_producer is not None and not app.in_transaction:
             if not self.queue:
                 return 0
             # ``Queue._queue`` is a CPython implementation detail (the
@@ -180,7 +182,7 @@ def __init__(
         assert api_version is not None
         super().__init__(loop=loop, **kwargs)
         self.buffer = ProducerBuffer(loop=self.loop, beacon=self.beacon)
-        if conf.producer_threaded:
+        if conf.producer_threaded and not self.app.in_transaction:
             # XXX ``create_threaded_producer`` is not defined here or on
             # ``ProducerT``; only the aiokafka driver implements it.  With the
             # confluent driver (faust/transport/drivers/confluent.py) this line

tests/unit/test_topics.py

@@ -162,6 +162,18 @@ async def test_publish_message__wait_enabled(self, *, topic, app):
         )
         callback.assert_called_once_with(fm)
 
+    @pytest.mark.asyncio
+    async def test_publish_message_uses_transaction_aware_producer(self, *, topic, app):
+        send_result = asyncio.get_running_loop().create_future()
+        producer = Mock(send=AsyncMock(return_value=send_result))
+        app.maybe_start_producer = AsyncMock(return_value=producer)
+        fm = topic.as_future_message(key="foo", value="bar")
+
+        await topic.publish_message(fm)
+
+        app.maybe_start_producer.assert_awaited_once_with()
+        producer.send.assert_awaited_once()
+
     @pytest.mark.asyncio
     async def test_publish_message__non_bytes_key_does_not_crash(self, *, topic, app):
         # A raw, unserialized non-bytes key/value can reach publish_message

tests/unit/transport/drivers/test_aiokafka.py

@@ -1436,6 +1436,12 @@ def assert_new_producer(
 
 
 class TestProducer(ProducerBaseTest):
+    def test_producer__threaded_disabled_for_transactions(self, *, app):
+        app.conf.producer_threaded = True
+        app.in_transaction = True
+        producer = Producer(app.transport)
+        assert producer.threaded_producer is None
+
     @pytest.mark.conf(producer_partitioner=my_partitioner)
     def test_producer__uses_custom_partitioner(self, *, producer):
         assert producer.partitioner is my_partitioner

tests/unit/transport/test_producer.py

@@ -1,6 +1,6 @@
 import asyncio
 from typing import Any, Optional
-from unittest.mock import Mock, PropertyMock, call
+from unittest.mock import Mock, PropertyMock, call, patch
 
 import pytest
 
@@ -22,6 +22,32 @@ def test_put(self, *, buf):
 
         buf.pending.put_nowait.assert_called_once_with(fut)
 
+    def test_put__threaded_when_configured(self, *, buf):
+        fut = Mock(name="future_message")
+        buf.app.conf.producer_threaded = True
+        buf.pending = Mock(name="pending")
+        threaded_producer = Mock(name="threaded_producer", thread_loop=Mock())
+        buf.threaded_producer = threaded_producer
+        buf.queue = Mock(name="queue")
+
+        with patch("faust.transport.producer.asyncio.run_coroutine_threadsafe") as run:
+            buf.put(fut)
+
+        run.assert_called_once()
+        buf.pending.put_nowait.assert_not_called()
+
+    def test_put__transactional_uses_async_buffer(self, *, buf):
+        fut = Mock(name="future_message")
+        buf.app.conf.producer_threaded = True
+        buf.app.in_transaction = True
+        buf.pending = Mock(name="pending")
+        buf.threaded_producer = Mock(name="threaded_producer")
+
+        buf.put(fut)
+
+        buf.pending.put_nowait.assert_called_once_with(fut)
+        buf.threaded_producer.event_queue.put.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_on_stop(self, *, buf):
         buf.flush = AsyncMock(name="flush")