feat: add consumer_member_id to metrics by default (#462)

* feat: add consumer_member_id to metrics by default

- Add ConsumerMetricsWrapper to automatically include consumer_member_id in all metrics
- Add get_consumer_metrics() function for creating consumer-aware metrics instances
- Update MetricsBuffer to use consumer-aware metrics automatically
- Update all processor metrics calls to include consumer member ID
- Add comprehensive tests for new functionality

This enables better observability and debugging in distributed consumer environments
by automatically tagging all consumer-related metrics with the consumer member ID.

* fix: resolve mypy strict type checking issues in consumer metrics

- Update MetricsBuffer record_* methods to use MetricName instead of str
- Add MetricName import to processor.py
- Fix test to use valid metric names from MetricName literal
- Ensure strict type safety compliance for consumer metrics wrapper

All mypy --strict checks now pass with no errors.

* Remove unnecessary metrics wrapper methods in MetricsBuffer

The record_timing, record_increment, and record_gauge wrapper methods
in MetricsBuffer claimed to add default tags but were just pass-through
functions that added no additional functionality. This removes the
misleading wrappers and replaces all usage with direct calls to the
underlying self.metrics methods.

Default tags like consumer member ID are already handled by
get_consumer_metrics(consumer.member_id) at the appropriate level.

* remove unused

---------

Co-authored-by: Markus Unterwaditzer <[email protected]>

Author: mwarkentin

arroyo/processing/processor.py

@@ -29,7 +29,7 @@
 )
 from arroyo.types import BrokerValue, Message, Partition, Topic, TStrategyPayload
 from arroyo.utils.logging import handle_internal_error
-from arroyo.utils.metrics import get_metrics
+from arroyo.utils.metrics import get_consumer_metrics
 
 logger = logging.getLogger(__name__)
 
@@ -89,8 +89,9 @@ class InvalidStateError(RuntimeError):
 
 
 class MetricsBuffer:
-    def __init__(self) -> None:
-        self.metrics = get_metrics()
+    def __init__(self, consumer: Consumer[Any]) -> None:
+        self.metrics = get_consumer_metrics(consumer.member_id)
+        self.__consumer = consumer
         self.__timers: MutableMapping[ConsumerTiming, float] = defaultdict(float)
         self.__counters: MutableMapping[ConsumerCounter, int] = defaultdict(int)
         self.__reset()
@@ -142,7 +143,7 @@ def __init__(
     ) -> None:
         self.__consumer = consumer
         self.__processor_factory = processor_factory
-        self.__metrics_buffer = MetricsBuffer()
+        self.__metrics_buffer = MetricsBuffer(consumer)
 
         self.__processing_strategy: Optional[
             ProcessingStrategy[TStrategyPayload]
@@ -236,7 +237,7 @@ 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.increment(
-                "arroyo.consumer.partitions_assigned.count", len(partitions), tags={"consumer_member_id": self.__consumer.member_id}
+                "arroyo.consumer.partitions_assigned.count", len(partitions)
             )
 
             current_partitions = dict(self.__consumer.tell())
@@ -262,7 +263,7 @@ def on_partitions_revoked(partitions: Sequence[Partition]) -> None:
             logger.info("Partitions to revoke: %r", partitions)
 
             self.__metrics_buffer.metrics.increment(
-                "arroyo.consumer.partitions_revoked.count", len(partitions), tags={"consumer_member_id": self.__consumer.member_id}
+                "arroyo.consumer.partitions_revoked.count", len(partitions)
             )
 
             if partitions:

arroyo/utils/metrics.py

@@ -45,6 +45,42 @@ def timing(
         raise NotImplementedError
 
 
+class ConsumerMetricsWrapper(Metrics):
+    """
+    A wrapper around a metrics backend that automatically adds consumer_member_id
+    to all metrics calls.
+    """
+
+    def __init__(self, metrics: Metrics, consumer_member_id: str) -> None:
+        self.__metrics = metrics
+        self.__consumer_member_id = 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
+
+    def increment(
+        self,
+        name: MetricName,
+        value: Union[int, float] = 1,
+        tags: Optional[Tags] = None,
+    ) -> None:
+        self.__metrics.increment(name, value, tags=self._add_consumer_tag(tags))
+
+    def gauge(
+        self, name: MetricName, value: Union[int, float], tags: Optional[Tags] = None
+    ) -> None:
+        self.__metrics.gauge(name, value, tags=self._add_consumer_tag(tags))
+
+    def timing(
+        self, name: MetricName, value: Union[int, float], tags: Optional[Tags] = None
+    ) -> None:
+        self.__metrics.timing(name, value, tags=self._add_consumer_tag(tags))
+
+
 class DummyMetricsBackend(Metrics):
     """
     Default metrics backend that does not record anything.
@@ -133,4 +169,12 @@ def get_metrics() -> Metrics:
     return _metrics_backend
 
 
-__all__ = ["configure_metrics", "Metrics", "MetricName"]
+def get_consumer_metrics(consumer_member_id: str) -> Metrics:
+    """
+    Get a metrics backend that automatically adds consumer_member_id to all metrics.
+    """
+    base_metrics = get_metrics()
+    return ConsumerMetricsWrapper(base_metrics, consumer_member_id)
+
+
+__all__ = ["configure_metrics", "Metrics", "MetricName", "Tags", "get_consumer_metrics"]

tests/utils/test_metrics.py

@@ -1,8 +1,8 @@
 import pytest
 
-from arroyo.utils.metrics import Gauge, MetricName, configure_metrics, get_metrics
+from arroyo.utils.metrics import Gauge, MetricName, configure_metrics, get_metrics, get_consumer_metrics
 from tests.metrics import Gauge as GaugeCall
-from tests.metrics import TestingMetricsBackend, _TestingMetricsBackend
+from tests.metrics import Increment, Timing, TestingMetricsBackend, _TestingMetricsBackend
 
 
 def test_gauge_simple() -> None:
@@ -31,3 +31,30 @@ def test_configure_metrics() -> None:
     # Can be reset to something else with force
     configure_metrics(_TestingMetricsBackend(), force=True)
     assert get_metrics() != TestingMetricsBackend
+
+
+def test_consumer_metrics_wrapper() -> None:
+    """Test that ConsumerMetricsWrapper automatically adds consumer_member_id to all metrics."""
+    # Reset to a fresh backend
+    backend = _TestingMetricsBackend()
+    configure_metrics(backend, force=True)
+
+    consumer_member_id = "test-consumer-123"
+    consumer_metrics = get_consumer_metrics(consumer_member_id)
+
+    # Test increment
+    consumer_metrics.increment("arroyo.consumer.run.count", 5, tags={"extra": "tag"})
+
+    # Test gauge
+    consumer_metrics.gauge("arroyo.consumer.librdkafka.total_queue_size", 10.5)
+
+    # Test timing
+    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"}),
+    ]
+
+    assert backend.calls == expected_calls