feat: add impact metric counter

Author: sjaanus

UnleashClient/impact_metrics/__init__.py

@@ -0,0 +1,19 @@
+from .metric_types import (
+    CollectedMetric,
+    Counter,
+    CounterImpl,
+    MetricLabels,
+    MetricOptions,
+    MetricType,
+    NumericMetricSample,
+)
+
+__all__ = [
+    "CollectedMetric",
+    "Counter",
+    "CounterImpl",
+    "MetricLabels",
+    "MetricOptions",
+    "MetricType",
+    "NumericMetricSample",
+]

UnleashClient/impact_metrics/metric_types.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from threading import RLock
+from typing import Any, Dict, List, Optional, Protocol
+
+
+class MetricType(str, Enum):
+    COUNTER = "counter"
+    GAUGE = "gauge"
+    HISTOGRAM = "histogram"
+
+
+MetricLabels = Dict[str, str]
+
+
+@dataclass
+class MetricOptions:
+    name: str
+    help: str
+    label_names: List[str] = field(default_factory=list)
+
+
+@dataclass
+class NumericMetricSample:
+    labels: MetricLabels
+    value: int
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"labels": self.labels, "value": self.value}
+
+
+@dataclass
+class CollectedMetric:
+    name: str
+    help: str
+    type: MetricType
+    samples: List[NumericMetricSample]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "name": self.name,
+            "help": self.help,
+            "type": self.type.value,
+            "samples": [s.to_dict() for s in self.samples],
+        }
+
+
+def _get_label_key(labels: Optional[MetricLabels]) -> str:
+    if not labels:
+        return ""
+    return ",".join(f"{k}={v}" for k, v in sorted(labels.items()))
+
+
+def _parse_label_key(key: str) -> MetricLabels:
+    if not key:
+        return {}
+    labels: MetricLabels = {}
+    for pair in key.split(","):
+        if "=" in pair:
+            k, v = pair.split("=", 1)
+            labels[k] = v
+    return labels
+
+
+class Counter(Protocol):
+    def inc(self, value: int = 1, labels: Optional[MetricLabels] = None) -> None: ...
+
+
+class CounterImpl:
+    def __init__(self, opts: MetricOptions) -> None:
+        self._opts = opts
+        self._values: Dict[str, int] = {}
+        self._lock = RLock()
+
+    def inc(self, value: int = 1, labels: Optional[MetricLabels] = None) -> None:
+        key = _get_label_key(labels)
+        with self._lock:
+            current = self._values.get(key, 0)
+            self._values[key] = current + value
+
+    def collect(self) -> CollectedMetric:
+        samples: List[NumericMetricSample] = []
+
+        with self._lock:
+            for key in list(self._values.keys()):
+                value = self._values.pop(key)
+                samples.append(
+                    NumericMetricSample(labels=_parse_label_key(key), value=value)
+                )
+
+        if not samples:
+            samples.append(NumericMetricSample(labels={}, value=0))
+
+        return CollectedMetric(
+            name=self._opts.name,
+            help=self._opts.help,
+            type=MetricType.COUNTER,
+            samples=samples,
+        )

tests/unit_tests/impact_metrics/__init__.py

@@ -0,0 +1,76 @@
+from UnleashClient.impact_metrics.metric_types import (
+    CounterImpl,
+    MetricOptions,
+)
+
+
+def test_counter_increments_by_default_value():
+    counter = CounterImpl(MetricOptions(name="test_counter", help="testing"))
+
+    counter.inc()
+
+    result = counter.collect()
+
+    assert result.to_dict() == {
+        "name": "test_counter",
+        "help": "testing",
+        "type": "counter",
+        "samples": [{"labels": {}, "value": 1}],
+    }
+
+
+def test_counter_increments_with_custom_value_and_labels():
+    counter = CounterImpl(MetricOptions(name="labeled_counter", help="with labels"))
+
+    counter.inc(3, {"foo": "bar"})
+    counter.inc(2, {"foo": "bar"})
+
+    result = counter.collect()
+
+    assert result.to_dict() == {
+        "name": "labeled_counter",
+        "help": "with labels",
+        "type": "counter",
+        "samples": [{"labels": {"foo": "bar"}, "value": 5}],
+    }
+
+
+def test_different_label_combinations_are_stored_separately():
+    counter = CounterImpl(MetricOptions(name="multi_label", help="label test"))
+
+    counter.inc(1, {"a": "x"})
+    counter.inc(2, {"b": "y"})
+    counter.inc(3)
+
+    result = counter.collect()
+
+    result.samples.sort(key=lambda s: s.value)
+
+    assert result.to_dict() == {
+        "name": "multi_label",
+        "help": "label test",
+        "type": "counter",
+        "samples": [
+            {"labels": {"a": "x"}, "value": 1},
+            {"labels": {"b": "y"}, "value": 2},
+            {"labels": {}, "value": 3},
+        ],
+    }
+
+
+def test_collect_returns_counter_with_zero_value_after_flushing_previous_values():
+    counter = CounterImpl(MetricOptions(name="flush_test", help="flush"))
+
+    counter.inc(1)
+    first = counter.collect()
+    assert first is not None
+    assert len(first.samples) == 1
+
+    second = counter.collect()
+
+    assert second.to_dict() == {
+        "name": "flush_test",
+        "help": "flush",
+        "type": "counter",
+        "samples": [{"labels": {}, "value": 0}],
+    }