sentry-arroyo 2.25.0__py3-none-any.whl → 2.26.0__py3-none-any.whl

arroyo/processing/processor.py

@@ -89,9 +89,8 @@ ConsumerCounter = Literal[
 
 
 class MetricsBuffer:
-    def __init__(self, consumer: Consumer[Any]) -> None:
-        self.metrics = get_consumer_metrics(consumer.member_id)
-        self.__consumer = consumer
+    def __init__(self) -> None:
+        self.metrics = get_consumer_metrics()
         self.__timers: MutableMapping[ConsumerTiming, float] = defaultdict(float)
         self.__counters: MutableMapping[ConsumerCounter, int] = defaultdict(int)
         self.__reset()
@@ -140,10 +139,11 @@ class StreamProcessor(Generic[TStrategyPayload]):
         commit_policy: CommitPolicy = ONCE_PER_SECOND,
         dlq_policy: Optional[DlqPolicy[TStrategyPayload]] = None,
         join_timeout: Optional[float] = None,
+        shutdown_strategy_before_consumer: bool = False,
     ) -> None:
         self.__consumer = consumer
         self.__processor_factory = processor_factory
-        self.__metrics_buffer = MetricsBuffer(consumer)
+        self.__metrics_buffer = MetricsBuffer()
 
         self.__processing_strategy: Optional[
             ProcessingStrategy[TStrategyPayload]
@@ -159,6 +159,7 @@ class StreamProcessor(Generic[TStrategyPayload]):
         self.__commit_policy_state = commit_policy.get_state_machine()
         self.__join_timeout = join_timeout
         self.__shutdown_requested = False
+        self.__shutdown_strategy_before_consumer = shutdown_strategy_before_consumer
 
         # 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.
@@ -171,49 +172,7 @@ class StreamProcessor(Generic[TStrategyPayload]):
         )
 
         def _close_strategy() -> None:
-            start_close = time.time()
-
-            if self.__processing_strategy is None:
-                # Partitions are revoked when the consumer is shutting down, at
-                # which point we already have closed the consumer.
-                return
-
-            logger.info("Closing %r...", self.__processing_strategy)
-            self.__processing_strategy.close()
-
-            logger.info("Waiting for %r to exit...", self.__processing_strategy)
-
-            while True:
-                start_join = time.time()
-
-                try:
-                    self.__processing_strategy.join(self.__join_timeout)
-                    self.__metrics_buffer.incr_timing(
-                        "arroyo.consumer.join.time", time.time() - start_join
-                    )
-                    break
-                except InvalidMessage as e:
-                    self.__metrics_buffer.incr_timing(
-                        "arroyo.consumer.join.time", time.time() - start_join
-                    )
-                    self._handle_invalid_message(e)
-
-            logger.info(
-                "%r exited successfully, releasing assignment.",
-                self.__processing_strategy,
-            )
-            self.__processing_strategy = None
-            self.__message = None  # avoid leaking buffered messages across assignments
-            self.__is_paused = False
-            self._clear_backpressure()
-
-            value = time.time() - start_close
-
-            self.__metrics_buffer.metrics.timing(
-                "arroyo.consumer.run.close_strategy", value
-            )
-
-            self.__metrics_buffer.incr_timing("arroyo.consumer.shutdown.time", value)
+            self._close_processing_strategy()
 
         def _create_strategy(partitions: Mapping[Partition, int]) -> None:
             start_create = time.time()
@@ -236,6 +195,8 @@ class StreamProcessor(Generic[TStrategyPayload]):
         def on_partitions_assigned(partitions: Mapping[Partition, int]) -> None:
             logger.info("New partitions assigned: %r", partitions)
             logger.info("Member id: %r", self.__consumer.member_id)
+            self.__metrics_buffer.metrics.consumer_member_id = self.__consumer.member_id
+
             self.__metrics_buffer.metrics.increment(
                 "arroyo.consumer.partitions_assigned.count", len(partitions)
             )
@@ -245,6 +206,7 @@ class StreamProcessor(Generic[TStrategyPayload]):
 
             if self.__dlq_policy:
                 self.__dlq_policy.reset_dlq_limits(current_partitions)
+
             if current_partitions:
                 if self.__processing_strategy is not None:
                     # TODO: for cooperative-sticky rebalancing this can happen
@@ -300,6 +262,48 @@ class StreamProcessor(Generic[TStrategyPayload]):
             [topic], on_assign=on_partitions_assigned, on_revoke=on_partitions_revoked
         )
 
+    def _close_processing_strategy(self) -> None:
+        """Close the processing strategy and wait for it to exit."""
+        start_close = time.time()
+
+        if self.__processing_strategy is None:
+            # Partitions are revoked when the consumer is shutting down, at
+            # which point we already have closed the consumer.
+            return
+
+        logger.info("Closing %r...", self.__processing_strategy)
+        logger.info("Member id: %r", self.__consumer.member_id)
+        self.__processing_strategy.close()
+
+        logger.info("Waiting for %r to exit...", self.__processing_strategy)
+
+        while True:
+            start_join = time.time()
+
+            try:
+                self.__processing_strategy.join(self.__join_timeout)
+                self.__metrics_buffer.incr_timing(
+                    "arroyo.consumer.join.time", time.time() - start_join
+                )
+                break
+            except InvalidMessage as e:
+                self.__metrics_buffer.incr_timing(
+                    "arroyo.consumer.join.time", time.time() - start_join
+                )
+                self._handle_invalid_message(e)
+
+        logger.info("%r exited successfully", self.__processing_strategy)
+        self.__processing_strategy = None
+        self.__message = None
+        self.__is_paused = False
+        self._clear_backpressure()
+
+        value = time.time() - start_close
+        self.__metrics_buffer.metrics.timing(
+            "arroyo.consumer.run.close_strategy", value
+        )
+        self.__metrics_buffer.incr_timing("arroyo.consumer.shutdown.time", value)
+
     def __commit(self, offsets: Mapping[Partition, int], force: bool = False) -> None:
         """
         If force is passed, commit immediately and do not throttle. This should
@@ -516,6 +520,13 @@ class StreamProcessor(Generic[TStrategyPayload]):
         self.__shutdown_requested = True
 
     def _shutdown(self) -> None:
+        # If shutdown_strategy_before_consumer is set, work around an issue
+        # where rdkafka would revoke our partition, but then also immediately
+        # revoke our member ID as well, causing join() of the CommitStrategy
+        # (that is running in the partition revocation callback) to crash.
+        if self.__shutdown_strategy_before_consumer:
+            self._close_processing_strategy()
+
         # close the consumer
         logger.info("Stopping consumer")
         self.__metrics_buffer.flush()

arroyo/utils/metrics.py

@@ -49,18 +49,21 @@ class ConsumerMetricsWrapper(Metrics):
     """
     A wrapper around a metrics backend that automatically adds consumer_member_id
     to all metrics calls.
+
+    Right now we only use this to add tags to the metrics emitted by
+    StreamProcessor, but ideally all metrics, even those emitted by strategies
+    and application code, would get this tag. The metrics abstraction in arroyo
+    is not sufficient for this. We'd have to add a "add_global_tags" method
+    (similar to the concept of global tags in sentry) and users would have to
+    implement it.
     """
 
-    def __init__(self, metrics: Metrics, consumer_member_id: str) -> None:
+    def __init__(self, metrics: Metrics) -> None:
         self.__metrics = metrics
-        self.__consumer_member_id = consumer_member_id
+        self.consumer_member_id = ""
 
     def _add_consumer_tag(self, tags: Optional[Tags]) -> Tags:
-        """Add consumer_member_id to the provided tags."""
-        consumer_tags = {"consumer_member_id": self.__consumer_member_id}
-        if tags:
-            return {**consumer_tags, **tags}
-        return consumer_tags
+        return {**(tags or {}), "consumer_member_id": self.consumer_member_id}
 
     def increment(
         self,
@@ -169,12 +172,12 @@ def get_metrics() -> Metrics:
     return _metrics_backend
 
 
-def get_consumer_metrics(consumer_member_id: str) -> Metrics:
+def get_consumer_metrics() -> ConsumerMetricsWrapper:
     """
     Get a metrics backend that automatically adds consumer_member_id to all metrics.
     """
     base_metrics = get_metrics()
-    return ConsumerMetricsWrapper(base_metrics, consumer_member_id)
+    return ConsumerMetricsWrapper(base_metrics)
 
 
 __all__ = ["configure_metrics", "Metrics", "MetricName", "Tags", "get_consumer_metrics"]

{sentry_arroyo-2.25.0.dist-info → sentry_arroyo-2.26.0.dist-info}/METADATA

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

{sentry_arroyo-2.25.0.dist-info → sentry_arroyo-2.26.0.dist-info}/RECORD

@@ -16,7 +16,7 @@ arroyo/backends/local/storages/__init__.py,sha256=AGYujdAAcn3osoj9jq84IzTywYbkID
 arroyo/backends/local/storages/abstract.py,sha256=1qVQp6roxHkK6XT2aklZyZk1qq7RzcPN6Db_CA5--kg,2901
 arroyo/backends/local/storages/memory.py,sha256=AoKDsVZzBXkOJyWArKWp3vfGfU9xLlKFXE9gsJiMIzQ,2613
 arroyo/processing/__init__.py,sha256=vZVg0wJvJfoVzlzGvnL59bT6YNIRJNQ5t7oU045Qbk4,87
-arroyo/processing/processor.py,sha256=cx8bPE41ZPT4hEhGdwCNAOe2ZAQ9lxZlr82goNkFBZ0,20702
+arroyo/processing/processor.py,sha256=PddFijN7uJyNVq1XMBeOSlD1Z0L27kLBplH0X_HF4Kg,21235
 arroyo/processing/strategies/__init__.py,sha256=EU_JMb54eOxMxaC5mIFpI-sAF-X2ZScbE8czBZ7bQkY,1106
 arroyo/processing/strategies/abstract.py,sha256=nu7juEz_aQmQIH35Z8u--FBuLjkK8_LQ1hIG2xpw9AA,4808
 arroyo/processing/strategies/batching.py,sha256=s89xC6lQpBseEaApu1iNTipXGKeO95OMwinj2VBKn9s,4778
@@ -39,14 +39,14 @@ arroyo/utils/concurrent.py,sha256=dbdPinjqmxCQ7izUGFNbGjB3OxfSIO01bnCSTANaVOE,11
 arroyo/utils/logging.py,sha256=Y1PnhYcI9XNNEK0H13Ct2xKLr2Niuw0dxayc6sWnui8,606
 arroyo/utils/metricDefs.json,sha256=5LR4hiHwD9JLaKm8JcpshmVjKCmIZzCxwa5oJxRkzH0,10264
 arroyo/utils/metric_defs.py,sha256=hZ98tCimeW8W6Gpo_LVvYI4RGtAPRsM6kruf8TtvdEY,7283
-arroyo/utils/metrics.py,sha256=eQwfdsbMbhHbTzrUgO-z1Cg0viPwFqhEbcTN-RqsMZk,4836
+arroyo/utils/metrics.py,sha256=UucA2igsosNtjA2L8h8EOkWyjBQNai5H4Vqah6KcgN8,5021
 arroyo/utils/profiler.py,sha256=aiYy2RRPX_IiDIO7AnFM3hARaHCctS3rqUS5nrHXbSg,2452
 arroyo/utils/retries.py,sha256=4MRhHUR7da9x1ytlo7YETo8S9HEebXmPF2-mKP4xYz0,3445
 examples/transform_and_produce/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 examples/transform_and_produce/batched.py,sha256=st2R6qTneAtV0JFbKP30Ti3sJDYj8Jkbmta9JckKdZU,2636
 examples/transform_and_produce/script.py,sha256=8kSMIjQNqGYEVyE0PvrfJh-a_UYCrJSstTp_De7kyyg,2306
 examples/transform_and_produce/simple.py,sha256=H7xqxItjl4tx34wVW5dy6mB9G39QucAtxkJSBzVmjgA,1637
-sentry_arroyo-2.25.0.dist-info/licenses/LICENSE,sha256=0Ng3MFdEcnz0sVD1XvGBBzbavvNp_7OAM5yVObB46jU,10829
+sentry_arroyo-2.26.0.dist-info/licenses/LICENSE,sha256=0Ng3MFdEcnz0sVD1XvGBBzbavvNp_7OAM5yVObB46jU,10829
 tests/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/backends/mixins.py,sha256=sfNyE0VTeiD3GHOnBYl-9urvPuURI2G1BWke0cz7Dvc,20445
 tests/backends/test_commit.py,sha256=iTHfK1qsBxim0XwxgMvNNSMqDUMEHoYkYBDcgxGBFbs,831
@@ -71,9 +71,9 @@ tests/processing/strategies/test_run_task_with_multiprocessing.py,sha256=eyv3O5X
 tests/processing/strategies/test_unfold.py,sha256=mbC4XhT6GkJRuC7vPR0h7jqwt4cu20q7Z114EJ6J9mQ,2009
 tests/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/utils/test_concurrent.py,sha256=Gwdzym2UZ1HO3rhOSGmzxImWcLFygY8P7MXHT3Q0xTE,455
-tests/utils/test_metrics.py,sha256=CL9cfWhBFm63Z5svNFIiLaddKsL1bINcOeKiekziQbs,2050
+tests/utils/test_metrics.py,sha256=y54LVGIkXok5cNVKKrdCpHmpZ0uJQOefT-PLvqZzjcw,2280
 tests/utils/test_retries.py,sha256=AxJLkXWeL9AjHv_p1n0pe8CXXJp24ZQIuYBHfNcmiz4,3075
-sentry_arroyo-2.25.0.dist-info/METADATA,sha256=7i3y9LsMIaUw4QFDsMzVn0V09ISVaB9s29hYeHXVqiQ,2208
-sentry_arroyo-2.25.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sentry_arroyo-2.25.0.dist-info/top_level.txt,sha256=DVdMZKysL_iIxm5aY0sYgZtP5ZXMg9YBaBmGQHVmDXA,22
-sentry_arroyo-2.25.0.dist-info/RECORD,,
+sentry_arroyo-2.26.0.dist-info/METADATA,sha256=OiWrlIZUPspZgkXQ3UgNef4O8bRi3JfXDgh-Sm6etyk,2208
+sentry_arroyo-2.26.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sentry_arroyo-2.26.0.dist-info/top_level.txt,sha256=DVdMZKysL_iIxm5aY0sYgZtP5ZXMg9YBaBmGQHVmDXA,22
+sentry_arroyo-2.26.0.dist-info/RECORD,,

tests/utils/test_metrics.py

@@ -1,8 +1,19 @@
 import pytest
 
-from arroyo.utils.metrics import Gauge, MetricName, configure_metrics, get_metrics, get_consumer_metrics
+from arroyo.utils.metrics import (
+    Gauge,
+    MetricName,
+    configure_metrics,
+    get_consumer_metrics,
+    get_metrics,
+)
 from tests.metrics import Gauge as GaugeCall
-from tests.metrics import Increment, Timing, TestingMetricsBackend, _TestingMetricsBackend
+from tests.metrics import (
+    Increment,
+    TestingMetricsBackend,
+    Timing,
+    _TestingMetricsBackend,
+)
 
 
 def test_gauge_simple() -> None:
@@ -40,7 +51,8 @@ def test_consumer_metrics_wrapper() -> None:
     configure_metrics(backend, force=True)
 
     consumer_member_id = "test-consumer-123"
-    consumer_metrics = get_consumer_metrics(consumer_member_id)
+    consumer_metrics = get_consumer_metrics()
+    consumer_metrics.consumer_member_id = consumer_member_id
 
     # Test increment
     consumer_metrics.increment("arroyo.consumer.run.count", 5, tags={"extra": "tag"})
@@ -52,9 +64,21 @@ def test_consumer_metrics_wrapper() -> None:
     consumer_metrics.timing("arroyo.consumer.poll.time", 100, tags={"another": "tag"})
 
     expected_calls = [
-        Increment("arroyo.consumer.run.count", 5, {"consumer_member_id": consumer_member_id, "extra": "tag"}),
-        GaugeCall("arroyo.consumer.librdkafka.total_queue_size", 10.5, {"consumer_member_id": consumer_member_id}),
-        Timing("arroyo.consumer.poll.time", 100, {"consumer_member_id": consumer_member_id, "another": "tag"}),
+        Increment(
+            "arroyo.consumer.run.count",
+            5,
+            {"consumer_member_id": consumer_member_id, "extra": "tag"},
+        ),
+        GaugeCall(
+            "arroyo.consumer.librdkafka.total_queue_size",
+            10.5,
+            {"consumer_member_id": consumer_member_id},
+        ),
+        Timing(
+            "arroyo.consumer.poll.time",
+            100,
+            {"consumer_member_id": consumer_member_id, "another": "tag"},
+        ),
     ]
 
     assert backend.calls == expected_calls