diff --git a/faust/transport/producer.py b/faust/transport/producer.py index f9967672f..49cb08478 100644 --- a/faust/transport/producer.py +++ b/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 diff --git a/tests/unit/test_topics.py b/tests/unit/test_topics.py index e880ce000..65338065e 100644 --- a/tests/unit/test_topics.py +++ b/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 diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index 21df98711..a9c52ebbe 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/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 diff --git a/tests/unit/transport/test_producer.py b/tests/unit/transport/test_producer.py index 0d1b6c3a3..15a7943b7 100644 --- a/tests/unit/transport/test_producer.py +++ b/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")