From 7a28935d270852e12027b4e60da7a3582686495e Mon Sep 17 00:00:00 2001 From: Lyn Date: Fri, 18 Aug 2023 10:37:01 -0700 Subject: [PATCH 1/2] fix(produce): Apply backpressure instead of crashing If we get local queue full, let's raise MessageRejected to slow down the consumer. --- arroyo/processing/strategies/produce.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/arroyo/processing/strategies/produce.py b/arroyo/processing/strategies/produce.py index 3b3b5a8a..cc9fef7d 100644 --- a/arroyo/processing/strategies/produce.py +++ b/arroyo/processing/strategies/produce.py @@ -87,7 +87,11 @@ def submit( future: Optional[Future[BrokerValue[TStrategyPayload]]] = None if not isinstance(message.payload, FilteredPayload): - future = self.__producer.produce(self.__topic, message.payload) + try: + future = self.__producer.produce(self.__topic, message.payload) + except BufferError as exc: + logger.exception(exc) + raise MessageRejected from exc self.__queue.append((message, future)) From 3c3ff4c5950caf6d6ead4fa75f38dd6d828b368b Mon Sep 17 00:00:00 2001 From: Lyn Date: Thu, 31 Aug 2023 07:18:43 -0700 Subject: [PATCH 2/2] add test --- tests/processing/strategies/test_produce.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/processing/strategies/test_produce.py b/tests/processing/strategies/test_produce.py index 212ac70e..27959ce9 100644 --- a/tests/processing/strategies/test_produce.py +++ b/tests/processing/strategies/test_produce.py @@ -1,8 +1,11 @@ from unittest import mock +import pytest + from arroyo.backends.kafka import KafkaPayload from arroyo.backends.local.backend import LocalBroker from arroyo.backends.local.storages.memory import MemoryMessageStorage +from arroyo.processing.strategies.abstract import MessageRejected from arroyo.processing.strategies.produce import Produce from arroyo.types import Message, Partition, Topic, Value from arroyo.utils.clock import TestingClock @@ -19,7 +22,7 @@ def test_produce() -> None: producer = broker.get_producer() next_step = mock.Mock() - strategy = Produce(producer, result_topic, next_step) + strategy = Produce(producer, result_topic, next_step, 2) value = b'{"something": "something"}' data = KafkaPayload(None, value, []) @@ -41,4 +44,10 @@ def test_produce() -> None: strategy.poll() assert next_step.submit.call_count == 2 assert next_step.poll.call_count == 2 + + # Backpressure if buffer size = 2 exceeded + with pytest.raises(MessageRejected): + for _ in range(3): + strategy.submit(message) + strategy.join()