fix(metrics): Re-set consumer member ID tag after every partition assignment (#466)

* fix(metrics): Re-set consumer member ID tag after every partition assignment

The member ID tag is set too early to a static value.

* fix tests

Author: untitaker

arroyo/processing/processor.py

@@ -89,9 +89,8 @@ class InvalidStateError(RuntimeError):
 
 
 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()
@@ -143,7 +142,7 @@ def __init__(
     ) -> None:
         self.__consumer = consumer
         self.__processor_factory = processor_factory
-        self.__metrics_buffer = MetricsBuffer(consumer)
+        self.__metrics_buffer = MetricsBuffer()
 
         self.__processing_strategy: Optional[
             ProcessingStrategy[TStrategyPayload]
@@ -236,6 +235,8 @@ def _create_strategy(partitions: Mapping[Partition, int]) -> None:
         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 +246,7 @@ def on_partitions_assigned(partitions: Mapping[Partition, int]) -> None:
 
             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

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"]

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