Configure the producer flush timeout with the max.poll.interval.ms (#382)

Author: ovv

quixstreams/app.py

@@ -50,6 +50,8 @@
 # Enforce idempotent producing for the internal RowProducer
 _default_producer_extra_config = {"enable.idempotence": True}
 
+_default_max_poll_interval_ms = 300000
+
 
 class Application:
     """
@@ -260,6 +262,10 @@ def __init__(
             broker_address=broker_address,
             extra_config=producer_extra_config,
             on_error=on_producer_error,
+            flush_timeout=consumer_extra_config.get(
+                "max.poll.interval.ms", _default_max_poll_interval_ms
+            )
+            / 1000,  # convert to seconds
         )
         self._consumer_poll_timeout = consumer_poll_timeout
         self._producer_poll_timeout = producer_poll_timeout

quixstreams/checkpointing/checkpoint.py

@@ -12,7 +12,7 @@
     DEFAULT_STATE_STORE_NAME,
 )
 from quixstreams.state.exceptions import StoreTransactionFailed
-from .exceptions import InvalidStoredOffset
+from .exceptions import InvalidStoredOffset, CheckpointProducerTimeout
 
 logger = logging.getLogger(__name__)
 
@@ -129,7 +129,13 @@ def commit(self):
 
         # Step 2. Flush producer to trigger all delivery callbacks and ensure that
         # all messages are produced
-        self._producer.flush()
+        logger.debug("Checkpoint: flushing producer")
+        unproduced_msg_count = self._producer.flush()
+        if unproduced_msg_count > 0:
+            raise CheckpointProducerTimeout(
+                f"'{unproduced_msg_count}' messages failed to be produced before the producer flush timeout"
+            )
+
         # Get produced offsets after flushing the producer
         produced_offsets = self._producer.offsets
 
@@ -138,6 +144,7 @@ def commit(self):
             TopicPartition(topic=topic, partition=partition, offset=offset + 1)
             for (topic, partition), offset in self._tp_offsets.items()
         ]
+        logger.debug("Checkpoint: commiting consumer")
         self._consumer.commit(offsets=offsets, asynchronous=False)
 
         # Step 4. Flush state store partitions to the disk together with changelog

quixstreams/checkpointing/exceptions.py

@@ -2,3 +2,6 @@
 
 
 class InvalidStoredOffset(QuixException): ...
+
+
+class CheckpointProducerTimeout(QuixException): ...

quixstreams/kafka/producer.py

@@ -48,6 +48,7 @@ def __init__(
         logger: logging.Logger = logger,
         error_callback: Callable[[KafkaError], None] = _default_error_cb,
         extra_config: Optional[dict] = None,
+        flush_timeout: Optional[int] = None,
     ):
         """
         A wrapper around `confluent_kafka.Producer`.
@@ -64,6 +65,7 @@ def __init__(
         :param extra_config: A dictionary with additional options that
             will be passed to `confluent_kafka.Producer` as is.
             Note: values passed as arguments override values in `extra_config`.
+        :param flush_timeout: The time the producer is waiting for all messages to be delivered.
         """
         if isinstance(broker_address, str):
             broker_address = ConnectionConfig(bootstrap_servers=broker_address)
@@ -76,6 +78,7 @@ def __init__(
             **{"logger": logger, "error_cb": error_callback},
         }
         self._inner_producer: Optional[ConfluentProducer] = None
+        self._flush_timeout = flush_timeout or -1
 
     def produce(
         self,
@@ -151,11 +154,13 @@ def flush(self, timeout: Optional[float] = None) -> int:
         Wait for all messages in the Producer queue to be delivered.
 
         :param float timeout: time to attempt flushing (seconds).
-            None or -1 is infinite. Default: None
+            None use producer default or -1 is infinite. Default: None
 
         :return: number of messages remaining to flush
         """
-        return self._producer.flush(timeout=timeout if timeout is not None else -1)
+        return self._producer.flush(
+            timeout=timeout if timeout is not None else self._flush_timeout
+        )
 
     @property
     def _producer(self) -> ConfluentProducer:

quixstreams/rowproducer.py

@@ -31,17 +31,20 @@ class RowProducer:
         If producer fails and the callback returns `True`, the exception
         will be logged but not propagated.
         The default callback logs an exception and returns `False`.
+    :param flush_timeout: The time the producer is waiting for all messages to be delivered.
     """
 
     def __init__(
         self,
         broker_address: Union[str, ConnectionConfig],
         extra_config: dict = None,
         on_error: Optional[ProducerErrorCallback] = None,
+        flush_timeout: Optional[int] = None,
     ):
         self._producer = Producer(
             broker_address=broker_address,
             extra_config=extra_config,
+            flush_timeout=flush_timeout,
         )
 
         self._on_error: Optional[ProducerErrorCallback] = (

tests/test_quixstreams/test_checkpointing.py

@@ -6,6 +6,7 @@
 from confluent_kafka import TopicPartition
 
 from quixstreams.checkpointing import Checkpoint, InvalidStoredOffset
+from quixstreams.checkpointing.exceptions import CheckpointProducerTimeout
 from quixstreams.kafka import Consumer
 from quixstreams.rowproducer import RowProducer
 from quixstreams.state import StateStoreManager
@@ -31,6 +32,13 @@ def factory(
     return factory
 
 
+@pytest.fixture()
+def rowproducer_mock(request):
+    p = MagicMock(spec_set=RowProducer)
+    p.flush.return_value = getattr(request, "param", 0)
+    return p
+
+
 class TestCheckpoint:
     def test_empty_true(self, checkpoint_factory):
         checkpoint = checkpoint_factory()
@@ -68,13 +76,18 @@ def test_commit_no_state_success(
         assert tp.offset == processed_offset + 1
 
     def test_commit_with_state_no_changelog_success(
-        self, checkpoint_factory, consumer, state_manager_factory, topic_factory
+        self,
+        checkpoint_factory,
+        consumer,
+        state_manager_factory,
+        topic_factory,
+        rowproducer_mock,
     ):
         topic_name, _ = topic_factory()
-        producer_mock = MagicMock(spec_set=RowProducer)
-        state_manager = state_manager_factory(producer=producer_mock)
+
+        state_manager = state_manager_factory(producer=rowproducer_mock)
         checkpoint = checkpoint_factory(
-            consumer_=consumer, state_manager_=state_manager, producer_=producer_mock
+            consumer_=consumer, state_manager_=state_manager, producer_=rowproducer_mock
         )
         processed_offset = 999
         key, value, prefix = "key", "value", b"__key__"
@@ -95,7 +108,7 @@ def test_commit_with_state_no_changelog_success(
         assert tp.offset == processed_offset + 1
 
         # Check the producer is flushed
-        assert producer_mock.flush.call_count == 1
+        assert rowproducer_mock.flush.call_count == 1
 
         # Check the state is flushed
         assert tx.completed
@@ -186,34 +199,32 @@ def test_commit_with_state_and_changelog_no_updates_success(
         assert not store_partition.get_processed_offset()
 
     def test_commit_no_offsets_stored_noop(
-        self, checkpoint_factory, state_manager_factory, topic_factory
+        self, checkpoint_factory, state_manager_factory, topic_factory, rowproducer_mock
     ):
         topic_name, _ = topic_factory()
-        producer_mock = MagicMock(spec_set=RowProducer)
         consumer_mock = MagicMock(spec_set=Consumer)
-        state_manager = state_manager_factory(producer=producer_mock)
+        state_manager = state_manager_factory(producer=rowproducer_mock)
         checkpoint = checkpoint_factory(
             consumer_=consumer_mock,
             state_manager_=state_manager,
-            producer_=producer_mock,
+            producer_=rowproducer_mock,
         )
         # Commit the checkpoint without processing any messages
         checkpoint.commit()
 
         # Check nothing is committed
         assert not consumer_mock.commit.call_count
-        assert not producer_mock.flush.call_count
+        assert not rowproducer_mock.flush.call_count
 
     def test_commit_has_failed_transactions_fails(
-        self, checkpoint_factory, state_manager_factory, topic_factory
+        self, checkpoint_factory, state_manager_factory, topic_factory, rowproducer_mock
     ):
-        producer_mock = MagicMock(spec_set=RowProducer)
         consumer_mock = MagicMock(spec_set=Consumer)
-        state_manager = state_manager_factory(producer=producer_mock)
+        state_manager = state_manager_factory(producer=rowproducer_mock)
         checkpoint = checkpoint_factory(
             consumer_=consumer_mock,
             state_manager_=state_manager,
-            producer_=producer_mock,
+            producer_=rowproducer_mock,
         )
         processed_offset = 999
         key, value, prefix = "key", "value", b"__key__"
@@ -240,20 +251,19 @@ def test_commit_has_failed_transactions_fails(
             checkpoint.commit()
 
         # The producer should not flush
-        assert not producer_mock.flush.call_count
+        assert not rowproducer_mock.flush.call_count
         # Consumer should not commit
         assert not consumer_mock.commit.call_count
 
     def test_commit_producer_flush_fails(
-        self, checkpoint_factory, state_manager_factory, topic_factory
+        self, checkpoint_factory, state_manager_factory, topic_factory, rowproducer_mock
     ):
-        producer_mock = MagicMock(spec_set=RowProducer)
         consumer_mock = MagicMock(spec_set=Consumer)
-        state_manager = state_manager_factory(producer=producer_mock)
+        state_manager = state_manager_factory(producer=rowproducer_mock)
         checkpoint = checkpoint_factory(
             consumer_=consumer_mock,
             state_manager_=state_manager,
-            producer_=producer_mock,
+            producer_=rowproducer_mock,
         )
         processed_offset = 999
         key, value, prefix = "key", "value", b"__key__"
@@ -266,7 +276,7 @@ def test_commit_producer_flush_fails(
         tx.set(key=key, value=value, prefix=prefix)
         checkpoint.store_offset("topic", 0, processed_offset)
 
-        producer_mock.flush.side_effect = ValueError("Flush failure")
+        rowproducer_mock.flush.side_effect = ValueError("Flush failure")
         # Checkpoint commit should fail if producer failed to flush
         with pytest.raises(ValueError):
             checkpoint.commit()
@@ -278,15 +288,14 @@ def test_commit_producer_flush_fails(
         assert not tx.completed
 
     def test_commit_consumer_commit_fails(
-        self, checkpoint_factory, state_manager_factory, topic_factory
+        self, checkpoint_factory, state_manager_factory, topic_factory, rowproducer_mock
     ):
-        producer_mock = MagicMock(spec_set=RowProducer)
         consumer_mock = MagicMock(spec_set=Consumer)
-        state_manager = state_manager_factory(producer=producer_mock)
+        state_manager = state_manager_factory(producer=rowproducer_mock)
         checkpoint = checkpoint_factory(
             consumer_=consumer_mock,
             state_manager_=state_manager,
-            producer_=producer_mock,
+            producer_=rowproducer_mock,
         )
         processed_offset = 999
         key, value, prefix = "key", "value", b"__key__"
@@ -305,7 +314,7 @@ def test_commit_consumer_commit_fails(
             checkpoint.commit()
 
         # Producer should flush
-        assert producer_mock.flush.call_count
+        assert rowproducer_mock.flush.call_count
         # The transaction should remain prepared, but not completed
         assert tx.prepared
         assert not tx.completed
@@ -326,3 +335,22 @@ def test_get_store_transaction_success(self, checkpoint_factory, state_manager):
         assert tx
         tx2 = checkpoint.get_store_transaction("topic", 0, "default")
         assert tx2 is tx
+
+    @pytest.mark.parametrize("rowproducer_mock", [1], indirect=True)
+    def test_incomplete_flush(
+        self, checkpoint_factory, consumer, state_manager_factory, rowproducer_mock
+    ):
+
+        state_manager = state_manager_factory(producer=rowproducer_mock)
+        checkpoint = checkpoint_factory(
+            consumer_=consumer, state_manager_=state_manager, producer_=rowproducer_mock
+        )
+        checkpoint.store_offset("topic", 0, 0)
+
+        with pytest.raises(CheckpointProducerTimeout) as err:
+            checkpoint.commit()
+
+        assert (
+            str(err.value)
+            == "'1' messages failed to be produced before the producer flush timeout"
+        )