Merge PausingManager into RowConsumer (#853)

Co-authored-by: Remy Gwaramadze <gwaramadze@users.noreply.github.com>

Author: daniil-quix

quixstreams/app.py

@@ -49,7 +49,7 @@
     check_state_management_enabled,
     is_quix_deployment,
 )
-from .processing import PausingManager, ProcessingContext
+from .processing import ProcessingContext
 from .rowconsumer import RowConsumer
 from .rowproducer import RowProducer
 from .runtracker import RunTracker
@@ -349,9 +349,6 @@ def __init__(
 
         self._source_manager = SourceManager()
         self._sink_manager = SinkManager()
-        self._pausing_manager = PausingManager(
-            consumer=self._consumer, topic_manager=self._topic_manager
-        )
         self._dataframe_registry = DataFrameRegistry()
         self._processing_context = ProcessingContext(
             commit_interval=self._config.commit_interval,
@@ -361,7 +358,6 @@ def __init__(
             state_manager=self._state_manager,
             exactly_once=self._config.exactly_once,
             sink_manager=self._sink_manager,
-            pausing_manager=self._pausing_manager,
             dataframe_registry=self._dataframe_registry,
         )
         self._run_tracker = RunTracker(processing_context=self._processing_context)
@@ -876,7 +872,7 @@ def _run_dataframe(self):
             else:
                 process_message(dataframes_composed)
                 processing_context.commit_checkpoint()
-                processing_context.resume_ready_partitions()
+                self._consumer.resume_backpressured()
                 source_manager.raise_for_error()
                 printer.print()
                 run_tracker.update_status()
@@ -1029,7 +1025,7 @@ def _on_revoke(self, _, topic_partitions: List[TopicPartition]):
             self._processing_context.commit_checkpoint(force=True)
 
         self._revoke_state_partitions(topic_partitions=topic_partitions)
-        self._processing_context.on_partition_revoke()
+        self._consumer.reset_backpressure()
 
     def _on_lost(self, _, topic_partitions: List[TopicPartition]):
         """
@@ -1038,7 +1034,7 @@ def _on_lost(self, _, topic_partitions: List[TopicPartition]):
         logger.debug("Rebalancing: dropping lost partitions")
 
         self._revoke_state_partitions(topic_partitions=topic_partitions)
-        self._processing_context.on_partition_revoke()
+        self._consumer.reset_backpressure()
 
     def _revoke_state_partitions(self, topic_partitions: List[TopicPartition]):
         non_changelog_topics = self._topic_manager.non_changelog_topics

quixstreams/checkpointing/checkpoint.py

@@ -6,8 +6,7 @@
 from confluent_kafka import KafkaException, TopicPartition
 
 from quixstreams.dataframe import DataFrameRegistry
-from quixstreams.kafka import BaseConsumer
-from quixstreams.processing.pausing import PausingManager
+from quixstreams.rowconsumer import RowConsumer
 from quixstreams.rowproducer import RowProducer
 from quixstreams.sinks import SinkManager
 from quixstreams.sinks.base import SinkBackpressureError
@@ -124,10 +123,9 @@ def __init__(
         self,
         commit_interval: float,
         producer: RowProducer,
-        consumer: BaseConsumer,
+        consumer: RowConsumer,
         state_manager: StateStoreManager,
         sink_manager: SinkManager,
-        pausing_manager: PausingManager,
         dataframe_registry: DataFrameRegistry,
         exactly_once: bool = False,
         commit_every: int = 0,
@@ -141,7 +139,6 @@ def __init__(
         self._consumer = consumer
         self._producer = producer
         self._sink_manager = sink_manager
-        self._pausing_manager = pausing_manager
         self._dataframe_registry = dataframe_registry
         self._exactly_once = exactly_once
         if self._exactly_once:
@@ -216,7 +213,7 @@ def commit(self):
                 # Pause the assignment to let it cool down and seek it back to
                 # the first processed offsets of this Checkpoint (it must be equal
                 # to the last committed offset).
-                self._pausing_manager.pause(
+                self._consumer.trigger_backpressure(
                     resume_after=exc.retry_after,
                     offsets_to_seek=self._starting_tp_offsets.copy(),
                 )

quixstreams/kafka/consumer.py

@@ -27,6 +27,7 @@
     "Consumer",
     "AutoOffsetReset",
     "RebalancingCallback",
+    "raise_for_msg_error",
 )
 
 RebalancingCallback = Callable[[ConfluentConsumer, List[TopicPartition]], None]

quixstreams/processing/__init__.py

@@ -1,2 +1 @@
 from .context import ProcessingContext as ProcessingContext
-from .pausing import PausingManager as PausingManager

quixstreams/processing/context.py

@@ -6,7 +6,6 @@
 from quixstreams.checkpointing import Checkpoint
 from quixstreams.dataframe import DataFrameRegistry
 from quixstreams.exceptions import QuixException
-from quixstreams.processing.pausing import PausingManager
 from quixstreams.rowconsumer import RowConsumer
 from quixstreams.rowproducer import RowProducer
 from quixstreams.sinks import SinkManager
@@ -33,7 +32,6 @@ class ProcessingContext:
     consumer: RowConsumer
     state_manager: StateStoreManager
     sink_manager: SinkManager
-    pausing_manager: PausingManager
     dataframe_registry: DataFrameRegistry
     commit_every: int = 0
     exactly_once: bool = False
@@ -70,7 +68,6 @@ def init_checkpoint(self):
             producer=self.producer,
             consumer=self.consumer,
             sink_manager=self.sink_manager,
-            pausing_manager=self.pausing_manager,
             dataframe_registry=self.dataframe_registry,
             exactly_once=self.exactly_once,
         )
@@ -98,12 +95,6 @@ def commit_checkpoint(self, force: bool = False):
                 )
             self.init_checkpoint()
 
-    def resume_ready_partitions(self):
-        self.pausing_manager.resume_if_ready()
-
-    def on_partition_revoke(self):
-        self.pausing_manager.reset()
-
     def __enter__(self):
         self.sink_manager.start_sinks()
         return self

quixstreams/processing/pausing.py

@@ -1,5 +1,7 @@
 import logging
-from typing import Callable, List, Mapping, Optional, Union
+import sys
+from time import monotonic
+from typing import Callable, List, Optional, Union
 
 from confluent_kafka import KafkaError, TopicPartition
 
@@ -14,8 +16,12 @@
 
 __all__ = ("RowConsumer",)
 
+_MAX_FLOAT = sys.float_info.max
+
 
 class RowConsumer(BaseConsumer):
+    _backpressure_resume_at: float
+
     def __init__(
         self,
         broker_address: Union[str, ConnectionConfig],
@@ -29,6 +35,7 @@ def __init__(
         on_error: Optional[ConsumerErrorCallback] = None,
     ):
         """
+
         A consumer class that is capable of deserializing Kafka messages to Rows
         according to the Topics deserialization settings.
 
@@ -66,7 +73,9 @@ def __init__(
             extra_config=extra_config,
         )
         self._on_error: ConsumerErrorCallback = on_error or default_on_consumer_error
-        self._topics: Mapping[str, Topic] = {}
+        self._topics: dict[str, Topic] = {}
+        self._backpressurred_tps: set[TopicPartition] = set()
+        self.reset_backpressure()
 
     def subscribe(
         self,
@@ -149,4 +158,72 @@ def poll_row(self, timeout: Optional[float] = None) -> Union[Row, List[Row], Non
 
     def close(self):
         super().close()
+        self.reset_backpressure()
         self._inner_consumer = None
+
+    @property
+    def backpressured_tps(self) -> set[TopicPartition]:
+        return self._backpressurred_tps
+
+    def trigger_backpressure(
+        self,
+        offsets_to_seek: dict[tuple[str, int], int],
+        resume_after: float,
+    ):
+        """
+        Pause all partitions for the certain period of time and seek the partitions
+        provided in the `offsets_to_seek` dict.
+
+        This method is supposed to be called in case of backpressure from Sinks.
+        """
+        resume_at = monotonic() + resume_after
+        self._backpressure_resume_at = min(self._backpressure_resume_at, resume_at)
+
+        changelog_topics = {k for k, v in self._topics.items() if v.is_changelog}
+        for tp in self.assignment():
+            # Pause only data TPs excluding changelog TPs
+            if tp.topic in changelog_topics:
+                continue
+
+            position, *_ = self.position([tp])
+            logger.debug(
+                f'Pausing topic partition "{tp.topic}[{tp.partition}]" for {resume_after}s; '
+                f"position={position.offset}"
+            )
+            self.pause(partitions=[tp])
+
+            # Seek the TP back to the "offset_to_seek" to start from it on resume.
+            # The "offset_to_seek" is provided by the Checkpoint and is expected to be the
+            # first offset processed in the checkpoint.
+            # There may be no offset for the TP if no message has been processed yet.
+            seek_offset = offsets_to_seek.get((tp.topic, tp.partition))
+            if seek_offset is not None:
+                logger.debug(
+                    f'Seek the paused partition "{tp.topic}[{tp.partition}]" back to '
+                    f"offset {seek_offset}"
+                )
+                self.seek(
+                    partition=TopicPartition(
+                        topic=tp.topic, partition=tp.partition, offset=seek_offset
+                    )
+                )
+
+            self._backpressurred_tps.add(tp)
+
+    def resume_backpressured(self):
+        """
+        Resume consuming from assigned data partitions after the wait period has elapsed.
+        """
+        if self._backpressure_resume_at > monotonic():
+            return
+
+        # Resume the previously backpressured TPs
+        for tp in self._backpressurred_tps:
+            logger.debug(f'Resuming topic partition "{tp.topic}[{tp.partition}]"')
+            self.resume(partitions=[tp])
+        self.reset_backpressure()
+
+    def reset_backpressure(self):
+        # Reset the timeout back to its initial state
+        self._backpressure_resume_at = _MAX_FLOAT
+        self._backpressurred_tps.clear()