Fix MetricsExporter mixing up instruments and labels when object ids are reused

changelog/fixed/mislabeled-marmoset.md

@@ -0,0 +1 @@
+Fix `contrib.opentelemetry` `MetricsExporter` recording a metric under another metric's instrument or another attribute set's labels when instruments or attributes are created per call. It cached them by `id()`, which Python reuses once an object is freed.

temporalio/contrib/opentelemetry/_metrics_exporter.py

@@ -7,6 +7,7 @@
 from collections.abc import Callable, Iterable, Mapping
 from dataclasses import dataclass, field
 from datetime import timedelta
+from typing import Generic, TypeVar
 
 import opentelemetry.metrics
 import opentelemetry.util.types
@@ -16,18 +17,55 @@
 
 logger = logging.getLogger(__name__)
 
+_K = TypeVar("_K")
+_V = TypeVar("_V")
+
+
+class _IdCache(Generic[_K, _V]):
+    # Keyed by id() for speed. Each entry holds its object, so the id can't be
+    # reused by another object while the entry is cached. A full cache starts
+    # over, which bounds what it keeps alive.
+    def __init__(self, max_size: int) -> None:
+        self._max_size = max_size
+        self._entries: dict[int, tuple[_K, _V]] = {}
+
+    def get(self, obj: _K) -> _V | None:
+        entry = self._entries.get(id(obj))
+        return None if entry is None else entry[1]
+
+    def put(self, obj: _K, value: _V) -> None:
+        if len(self._entries) >= self._max_size:
+            self._entries.clear()
+        self._entries[id(obj)] = (obj, value)
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+
+_ID_CACHE_SIZE = 1000
+
 
 @dataclass
 class _GaugeState:
     lock: threading.Lock = field(default_factory=threading.Lock)
-    # Keyed by id(BufferedMetricUpdate.attributes) so repeated updates for the
-    # same attribute set overwrite rather than accumulate -- gauges are
-    # "latest value wins", unlike counters.
+    # Keyed by attribute values so repeated updates for the same attribute set
+    # overwrite rather than accumulate -- gauges are "latest value wins",
+    # unlike counters.
     last_values: dict[
-        int, tuple[Mapping[str, opentelemetry.util.types.AttributeValue], int | float]
+        frozenset[tuple[str, opentelemetry.util.types.AttributeValue]],
+        tuple[Mapping[str, opentelemetry.util.types.AttributeValue], int | float],
     ] = field(default_factory=dict)
 
 
+@dataclass(slots=True)
+class _ConvertedAttributes:
+    attributes: Mapping[str, opentelemetry.util.types.AttributeValue]
+    # Only gauges key on it, so it is built on the first gauge update.
+    gauge_key: frozenset[tuple[str, opentelemetry.util.types.AttributeValue]] | None = (
+        None
+    )
+
+
 @dataclass
 class _Instrument:
     kind: temporalio.runtime.BufferedMetricKind
@@ -136,10 +174,15 @@ def __init__(
         ).get_meter(meter_name, meter_version)
         self._poll_interval = poll_interval
         self._on_error = on_error
-        self._instruments: dict[int, _Instrument] = {}
-        self._attributes_cache: dict[
-            int, Mapping[str, opentelemetry.util.types.AttributeValue]
+        self._instruments: dict[
+            tuple[str, temporalio.runtime.BufferedMetricKind], _Instrument
         ] = {}
+        self._instruments_by_metric: _IdCache[
+            temporalio.runtime.BufferedMetric, _Instrument
+        ] = _IdCache(_ID_CACHE_SIZE)
+        self._attributes_cache: _IdCache[
+            temporalio.common.MetricAttributes, _ConvertedAttributes
+        ] = _IdCache(_ID_CACHE_SIZE)
         self._started = False
         self._shutdown_event = asyncio.Event()
         self._run_complete_event = asyncio.Event()
@@ -221,19 +264,21 @@ def _handle_error(self, err: Exception) -> None:
 
     def _apply_update(self, update: temporalio.runtime.BufferedMetricUpdate) -> None:
         instrument = self._get_or_create_instrument(update.metric)
-        attributes = self._get_or_create_attributes(update.attributes)
+        converted = self._get_or_create_attributes(update.attributes)
 
         if instrument.kind == temporalio.runtime.BUFFERED_METRIC_KIND_COUNTER:
             assert instrument.counter is not None
-            instrument.counter.add(update.value, attributes=attributes)
+            instrument.counter.add(update.value, attributes=converted.attributes)
         elif instrument.kind == temporalio.runtime.BUFFERED_METRIC_KIND_HISTOGRAM:
             assert instrument.histogram is not None
-            instrument.histogram.record(update.value, attributes=attributes)
+            instrument.histogram.record(update.value, attributes=converted.attributes)
         elif instrument.kind == temporalio.runtime.BUFFERED_METRIC_KIND_GAUGE:
             assert instrument.gauge_state is not None
+            if converted.gauge_key is None:
+                converted.gauge_key = frozenset(converted.attributes.items())
             with instrument.gauge_state.lock:
-                instrument.gauge_state.last_values[id(update.attributes)] = (
-                    attributes,
+                instrument.gauge_state.last_values[converted.gauge_key] = (
+                    converted.attributes,
                     update.value,
                 )
         else:
@@ -242,9 +287,14 @@ def _apply_update(self, update: temporalio.runtime.BufferedMetricUpdate) -> None
     def _get_or_create_instrument(
         self, metric: temporalio.runtime.BufferedMetric
     ) -> _Instrument:
-        instrument = self._instruments.get(id(metric))
+        instrument = self._instruments_by_metric.get(metric)
         if instrument is not None:
             return instrument
+        key = (metric.name, metric.kind)
+        instrument = self._instruments.get(key)
+        if instrument is not None:
+            self._instruments_by_metric.put(metric, instrument)
+            return instrument
 
         unit = metric.unit or ""
         description = metric.description or ""
@@ -269,17 +319,16 @@ def _get_or_create_instrument(
         else:
             raise ValueError(f"Unrecognized buffered metric kind: {metric.kind}")
 
-        self._instruments[id(metric)] = instrument
+        self._instruments[key] = instrument
+        self._instruments_by_metric.put(metric, instrument)
         return instrument
 
     def _get_or_create_attributes(
         self, attributes: temporalio.common.MetricAttributes
-    ) -> Mapping[str, opentelemetry.util.types.AttributeValue]:
-        cached = self._attributes_cache.get(id(attributes))
-        if cached is not None:
-            return cached
-        converted: Mapping[str, opentelemetry.util.types.AttributeValue] = dict(
-            attributes
-        )
-        self._attributes_cache[id(attributes)] = converted
+    ) -> _ConvertedAttributes:
+        converted = self._attributes_cache.get(attributes)
+        if converted is not None:
+            return converted
+        converted = _ConvertedAttributes(dict(attributes))
+        self._attributes_cache.put(attributes, converted)
         return converted

tests/contrib/opentelemetry/test_metrics_exporter.py

@@ -7,7 +7,7 @@
 from opentelemetry.sdk.metrics import MeterProvider
 from opentelemetry.sdk.metrics.export import InMemoryMetricReader
 
-from temporalio.contrib.opentelemetry import MetricsExporter
+from temporalio.contrib.opentelemetry import MetricsExporter, _metrics_exporter
 from temporalio.runtime import MetricBuffer, Runtime, TelemetryConfig
 
 
@@ -132,6 +132,62 @@ async def test_metrics_exporter_attributes_passthrough():
     assert dict(points[0].attributes) == {"foo": "bar", "baz": 123}
 
 
+async def test_metrics_exporter_instruments_created_per_call():
+    runtime, _, _, reader, exporter = _make_runtime_and_exporter()
+
+    # Like an activity that creates its instruments on each call: each call
+    # brings new metric and attribute objects, and a cache keyed by id() that
+    # doesn't hold them would see their ids reused.
+    calls = [("model-a", 100, 900), ("model-b", 7, 40), ("model-c", 55, 300)]
+    for model, tokens, latency_ms in calls:
+        meter = runtime.metric_meter.with_additional_attributes({"model": model})
+        meter.create_counter("tokens").add(tokens)
+        meter.create_histogram("latency").record(latency_ms)
+        meter.create_gauge("last-tokens").set(tokens)
+        exporter._drain_once()
+
+    metrics = _collect_metrics_by_name(reader)
+    tokens_by_model = {model: tokens for model, tokens, _ in calls}
+    assert {
+        p.attributes["model"]: p.value for p in metrics["tokens"].data.data_points
+    } == tokens_by_model
+    assert {
+        p.attributes["model"]: p.value for p in metrics["last-tokens"].data.data_points
+    } == tokens_by_model
+    assert {
+        p.attributes["model"]: (p.count, p.sum)
+        for p in metrics["latency"].data.data_points
+    } == {model: (1, latency_ms) for model, _, latency_ms in calls}
+
+
+async def test_metrics_exporter_caches_are_bounded(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(_metrics_exporter, "_ID_CACHE_SIZE", 2)
+    runtime, _, _, reader, exporter = _make_runtime_and_exporter()
+
+    meters = [
+        runtime.metric_meter.with_additional_attributes({"n": n}) for n in range(5)
+    ]
+    counters = [meter.create_counter("bounded-counter") for meter in meters]
+    gauges = [meter.create_gauge("bounded-gauge") for meter in meters]
+    for value in (1, 2):
+        for counter, gauge in zip(counters, gauges):
+            counter.add(value)
+            gauge.set(value)
+        exporter._drain_once()
+        assert len(exporter._attributes_cache) <= 2
+        assert len(exporter._instruments_by_metric) <= 2
+
+    # Every set was evicted before its second update, so this also checks
+    # that an evicted object is looked up again correctly.
+    metrics = _collect_metrics_by_name(reader)
+    assert {
+        p.attributes["n"]: p.value for p in metrics["bounded-counter"].data.data_points
+    } == {n: 3 for n in range(5)}
+    assert {
+        p.attributes["n"]: p.value for p in metrics["bounded-gauge"].data.data_points
+    } == {n: 2 for n in range(5)}
+
+
 async def test_metrics_exporter_buffer_not_attached_raises():
     unattached_buffer = MetricBuffer(10_000)
     reader = InMemoryMetricReader()