sentry-arroyo 2.32.4__tar.gz → 2.32.5__tar.gz

{sentry_arroyo-2.32.4/sentry_arroyo.egg-info → sentry_arroyo-2.32.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sentry-arroyo
-Version: 2.32.4
+Version: 2.32.5
 Summary: Arroyo is a Python library for working with streaming data.
 Home-page: https://github.com/getsentry/arroyo
 Author: Sentry

{sentry_arroyo-2.32.4 → sentry_arroyo-2.32.5}/arroyo/processing/processor.py

@@ -141,14 +141,15 @@ class StreamProcessor(Generic[TStrategyPayload]):
         dlq_policy: Optional[DlqPolicy[TStrategyPayload]] = None,
         join_timeout: Optional[float] = None,
         shutdown_strategy_before_consumer: bool = False,
+        handle_poll_while_paused: Optional[bool] = False,
     ) -> None:
         self.__consumer = consumer
         self.__processor_factory = processor_factory
         self.__metrics_buffer = MetricsBuffer()
 
-        self.__processing_strategy: Optional[ProcessingStrategy[TStrategyPayload]] = (
-            None
-        )
+        self.__processing_strategy: Optional[
+            ProcessingStrategy[TStrategyPayload]
+        ] = None
 
         self.__message: Optional[BrokerValue[TStrategyPayload]] = None
 
@@ -167,6 +168,8 @@ class StreamProcessor(Generic[TStrategyPayload]):
         self.__shutdown_requested = False
         self.__shutdown_strategy_before_consumer = shutdown_strategy_before_consumer
 
+        self.__handle_poll_while_paused = handle_poll_while_paused
+
         # Buffers messages for DLQ. Messages are added when they are submitted for processing and
         # removed once the commit callback is fired as they are guaranteed to be valid at that point.
         self.__buffered_messages: BufferedMessages[TStrategyPayload] = BufferedMessages(
@@ -464,40 +467,59 @@ class StreamProcessor(Generic[TStrategyPayload]):
                         self.__is_paused = True
 
                     elif self.__is_paused:
-                        paused_partitions = set(self.__consumer.paused())
-                        all_partitions = set(self.__consumer.tell())
-                        unpaused_partitions = (
-                            all_partitions - paused_partitions
-                        )
-                        if unpaused_partitions:
-                            logger.warning(
-                                "Processor in paused state while consumer is partially unpaused: %s, paused: %s",
-                                unpaused_partitions,
-                                paused_partitions,
-                            )
-                            self.__is_paused = False
-                            # unpause paused partitions... just in case a subset is paused
-                            self.__metrics_buffer.incr_counter(
-                                "arroyo.consumer.resume", 1
-                            )
-                            self.__consumer.resume([*paused_partitions])
-                        else:
-                            # A paused consumer should still poll periodically to avoid it's partitions
-                            # getting revoked by the broker after reaching the max.poll.interval.ms
-                            # Polling a paused consumer should never yield a message.
-                            logger.warning("consumer.tell() value right before poll() is: %s", self.__consumer.tell())
+                        if self.__handle_poll_while_paused:
                             maybe_message = self.__consumer.poll(0.1)
                             if maybe_message is not None:
-                                logger.warning("Received a message from partition: %s, \
-                                                consumer.tell() value right after poll() is: %s \
-                                                Some lines above consumer.tell() was called, all_partitons value was: %s \
-                                                Some lines above consumer.paused() was called, paused_partitions value is: %s",
-                                                maybe_message.partition,
-                                                self.__consumer.tell(),
-                                                all_partitions,
-                                                paused_partitions
-                                                )
-                            assert maybe_message is None
+                                # The paused consumer, in the above poll, has
+                                # gone through rebalancing. In this case we
+                                # expect that partition revocation cleared the
+                                # pause flag and the carried over message.
+                                # this assumption will not hold for cooperative-sticky rebalancing.
+                                assert (
+                                    not self.__is_paused
+                                ), "consumer unpaused itself without rebalancing"
+                                assert (
+                                    self.__message is None
+                                ), "consumer unpaused itself without rebalancing"
+                        # this path might raise AssertionErrors
+                        else:
+                            paused_partitions = set(self.__consumer.paused())
+                            all_partitions = set(self.__consumer.tell())
+                            unpaused_partitions = all_partitions - paused_partitions
+                            if unpaused_partitions:
+                                logger.warning(
+                                    "Processor in paused state while consumer is partially unpaused: %s, paused: %s",
+                                    unpaused_partitions,
+                                    paused_partitions,
+                                )
+                                self.__is_paused = False
+                                # unpause paused partitions... just in case a subset is paused
+                                self.__metrics_buffer.incr_counter(
+                                    "arroyo.consumer.resume", 1
+                                )
+                                self.__consumer.resume([*paused_partitions])
+                            else:
+                                # A paused consumer should still poll periodically to avoid it's partitions
+                                # getting revoked by the broker after reaching the max.poll.interval.ms
+                                # Polling a paused consumer should never yield a message.
+                                logger.warning(
+                                    "consumer.tell() value right before poll() is: %s",
+                                    self.__consumer.tell(),
+                                )
+                                maybe_message = self.__consumer.poll(0.1)
+                                if maybe_message is not None:
+                                    logger.warning(
+                                        "Received a message from partition: %s, \
+                                                    consumer.tell() value right after poll() is: %s \
+                                                    Some lines above consumer.tell() was called, all_partitons value was: %s \
+                                                    Some lines above consumer.paused() was called, paused_partitions value is: %s",
+                                        maybe_message.partition,
+                                        self.__consumer.tell(),
+                                        all_partitions,
+                                        paused_partitions,
+                                    )
+                                assert maybe_message is None
+
                     else:
                         time.sleep(0.01)
 

{sentry_arroyo-2.32.4 → sentry_arroyo-2.32.5/sentry_arroyo.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sentry-arroyo
-Version: 2.32.4
+Version: 2.32.5
 Summary: Arroyo is a Python library for working with streaming data.
 Home-page: https://github.com/getsentry/arroyo
 Author: Sentry

{sentry_arroyo-2.32.4 → sentry_arroyo-2.32.5}/setup.py

@@ -10,7 +10,7 @@ def get_requirements() -> Sequence[str]:
 
 setup(
     name="sentry-arroyo",
-    version="2.32.4",
+    version="2.32.5",
     author="Sentry",
     author_email="oss@sentry.io",
     license="Apache-2.0",

{sentry_arroyo-2.32.4 → sentry_arroyo-2.32.5}/tests/backends/test_kafka.py

@@ -80,7 +80,7 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
     @property
     def configuration(self) -> KafkaBrokerConfig:
         config = {
-            "bootstrap.servers": os.environ.get("DEFAULT_BROKERS", "localhost:9092"),
+            "bootstrap.servers": os.environ.get("DEFAULT_BROKERS", "127.0.0.1:9092"),
         }
 
         return build_kafka_configuration(config)
@@ -231,6 +231,93 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
                 with pytest.raises(RuntimeError):
                     processor.run()
 
+    @mock.patch("arroyo.processing.processor.BACKPRESSURE_THRESHOLD", 0)
+    def test_assign_partition_during_pause(self) -> None:
+        if self.cooperative_sticky or self.kip_848:
+            pytest.skip("test does not work with cooperative-sticky rebalancing")
+
+        payloads = self.get_payloads()
+
+        strategy = mock.Mock()
+        strategy.submit.side_effect = MessageRejected()
+        factory = mock.Mock()
+        factory.create_with_partitions.return_value = strategy
+
+        partition_count = 2
+
+        with self.get_topic(partition_count) as topic, closing(
+            self.get_producer()
+        ) as producer, closing(
+            self.get_consumer(
+                "test_assign_partition_during_pause", enable_end_of_partition=True
+            )
+        ) as consumer_a, closing(
+            self.get_consumer(
+                "test_assign_partition_during_pause", enable_end_of_partition=True
+            )
+        ) as consumer_b:
+            for i in range(partition_count):
+                producer.produce(Partition(topic, i), next(payloads)).result(
+                    timeout=5.0
+                )
+
+            processor_a = StreamProcessor(
+                consumer_a, topic, factory, IMMEDIATE, handle_poll_while_paused=True
+            )
+
+            def wait_until_consumer_pauses(processor: StreamProcessor[Any]) -> None:
+                for _ in range(20):
+                    try:
+                        processor._run_once()
+                    except EndOfPartition:
+                        pass
+
+                    if processor._StreamProcessor__is_paused:  # type:ignore
+                        return
+                raise RuntimeError("processor was not paused")
+
+            # calling _run_once will pause both consumers because of the MessageRejected strategy above
+            wait_until_consumer_pauses(processor_a)
+
+            # consumer A has all the partitions
+            assert len(consumer_a.tell()) == 2
+            assert len(consumer_b.tell()) == 0
+
+            # consumer A has all partitions paused (both from consumer and from
+            # StreamProcessor POV)
+            assert consumer_a.paused()
+            assert processor_a._StreamProcessor__is_paused is True  # type:ignore
+
+            # subscribe with another consumer, now we should have rebalancing in the next few polls
+            processor_b = StreamProcessor(
+                consumer_b, topic, factory, IMMEDIATE, handle_poll_while_paused=True
+            )
+
+            for _ in range(10):
+                try:
+                    processor_a._run_once()
+                except EndOfPartition:
+                    pass
+                try:
+                    processor_b._run_once()
+                except EndOfPartition:
+                    pass
+
+            # balanced
+            assert len(consumer_a.tell()) == 1
+            assert len(consumer_b.tell()) == 1
+
+            # close B, but A has not polled yet, so it only has one partition still
+            consumer_b.close()
+            assert len(consumer_a.tell()) == 1
+
+            for _ in range(20):
+                try:
+                    processor_a._run_once()
+                except EndOfPartition:
+                    pass
+            assert len(consumer_a.tell()) == 2
+
     def test_consumer_polls_when_paused(self) -> None:
         strategy = mock.Mock()
         factory = mock.Mock()

{sentry_arroyo-2.32.4 → sentry_arroyo-2.32.5}/tests/processing/test_processor.py

@@ -690,3 +690,75 @@ def test_processor_pause_with_invalid_message() -> None:
 
     processor._run_once()
     assert strategy.submit.call_args_list[-1] == mock.call(new_message)
+
+def test_processor_poll_while_paused() -> None:
+
+    topic = Topic("topic")
+
+    consumer = mock.Mock()
+    strategy = mock.Mock()
+    factory = mock.Mock()
+    factory.create_with_partitions.return_value = strategy
+
+    processor: StreamProcessor[int] = StreamProcessor(
+        consumer, topic, factory, IMMEDIATE, handle_poll_while_paused=True
+    )
+
+    # Subscribe to topic
+    subscribe_args, subscribe_kwargs = consumer.subscribe.call_args
+    assert subscribe_args[0] == [topic]
+
+    # Partition assignment
+    partition = Partition(topic, 0)
+    new_partition = Partition(topic, 1)
+    consumer.tell.return_value = {}
+    assignment_callback = subscribe_kwargs["on_assign"]
+    offsets = {partition: 0}
+    assignment_callback(offsets)
+
+    # Message that we will get from polling
+    message = Message(BrokerValue(0, partition, 0, datetime.now()))
+
+    # Message will be rejected
+    consumer.poll.return_value = message.value
+    strategy.submit.side_effect = MessageRejected()
+    with assert_changes(lambda: int(consumer.pause.call_count), 0, 1):
+        processor._run_once()
+        assert strategy.submit.call_args_list[-1] == mock.call(message)
+
+        with mock.patch("time.time", return_value=time.time() + 5):
+            processor._run_once()  # Should pause now
+
+    # Consumer is in paused state
+    # The same rejected message should be carried over
+
+    # All partitions are paused
+    consumer.paused.return_value = set(p for p in offsets)
+    # Simulate a continuous backpressure state where messages are being rejected
+    strategy.submit.side_effect = MessageRejected()
+
+    # Simulate Kafka returning nothing since the consumer is paused
+    consumer.poll.return_value = None
+
+    # The next poll returns nothing, but we are still carrying over the rejected message
+    processor._run_once()
+    assert consumer.poll.return_value is None
+
+    # At this point, let's say the message carried over is invalid (e.g. it could be stale)
+    strategy.submit.side_effect = InvalidMessage(partition, 0, needs_commit=False)
+
+    # Handles the invalid message and unpauses the consumer
+    with assert_changes(lambda: int(consumer.resume.call_count), 0, 1):
+        processor._run_once()
+
+    # Poll for the next message from Kafka, but this time the partition has changed
+    new_message = Message(BrokerValue(0, new_partition, 1, datetime.now()))
+    consumer.poll.return_value = new_message.value
+    processor._run_once()
+    assert processor._StreamProcessor__is_paused is False # type:ignore
+
+    strategy.submit.return_value = None
+    strategy.submit.side_effect = None
+
+    processor._run_once()
+    assert strategy.submit.call_args_list[-1] == mock.call(new_message)