feat(metrics): attach trace exemplars to flushed metric series (#1057)

Author: DanielVisca

.sampo/changesets/metrics-trace-exemplars.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+`client.metrics` now attaches a trace exemplar to each series it flushes. A sample recorded inside a sampled span (a PostHog `start_span` span, or the current OpenTelemetry span) stamps its data point with that span's trace ID, span ID, and value, so metric charts can link to a representative trace. One exemplar per series per flush window; the last sample wins.

posthog/metrics_capture.py

@@ -12,6 +12,11 @@
 is attached: every attribute value creates a new series, and per-user series
 are the canonical metrics-cardinality explosion.
 
+A sample recorded inside a sampled span (a PostHog ``start_span`` span, or the
+current OpenTelemetry span) stamps its series with an exemplar: the trace and
+span IDs plus the value. One exemplar per series per window, last sample wins,
+so a metric point links to one representative trace at no cardinality cost.
+
 Delivery is at-least-once: a request that succeeds server-side but fails
 client-side (e.g. a read timeout) is retried with the next window, which can
 double-count that window's deltas. Failed flushes retry with exponential
@@ -37,6 +42,8 @@
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
+from .tracing._traceparent import TRACE_FLAGS_SAMPLED, parse_traceparent
+
 log = logging.getLogger("posthog")
 
 MetricAttributeValue = Union[str, int, float, bool]
@@ -74,6 +81,8 @@
 _DEFAULT_FLUSH_INTERVAL_SECONDS = 10.0
 _DEFAULT_MAX_SERIES_PER_FLUSH = 1000
 _SCOPE_NAME = "posthog-python"
+# None until the first lookup; False when opentelemetry-api is not installed.
+_otel_trace: Any = None
 
 
 def _snapshot_attribute_value(value: Any) -> Any:
@@ -135,6 +144,52 @@ def _bucket_index_for(value: float, bounds: list) -> int:
     return len(bounds)
 
 
+def _parse_sampled_traceparent(value: Any) -> Optional[tuple[str, str]]:
+    """The (trace_id, span_id) of a sampled W3C ``traceparent``, or None."""
+    context = parse_traceparent(value)
+    if context is None or context.flags != TRACE_FLAGS_SAMPLED:
+        return None
+    return context.trace_id, context.span_id
+
+
+def _otel_trace_module() -> Any:
+    global _otel_trace
+    if _otel_trace is None:
+        try:
+            from opentelemetry import trace
+
+            _otel_trace = trace
+        except ImportError:
+            _otel_trace = False
+    return _otel_trace
+
+
+def _active_trace_context(client: Any) -> Optional[tuple[str, str]]:
+    """The sampled span a sample is recorded in: PostHog's own active span first,
+    then the current OpenTelemetry span. An unsampled span is never exported, so
+    an exemplar pointing at it would link to a trace that does not exist."""
+    try:
+        get_active_span = getattr(client, "get_active_span", None)
+        span = get_active_span() if callable(get_active_span) else None
+        if span is not None:
+            # A PostHog span is active, so it is the innermost context; an
+            # unsampled one must not fall through to an outer OTel span.
+            return _parse_sampled_traceparent(span.traceparent())
+    except Exception:
+        pass
+
+    trace = _otel_trace_module()
+    if not trace:
+        return None
+    try:
+        span_context = trace.get_current_span().get_span_context()
+        if not span_context.is_valid or not span_context.trace_flags.sampled:
+            return None
+        return f"{span_context.trace_id:032x}", f"{span_context.span_id:016x}"
+    except Exception:
+        return None
+
+
 def _series_key(
     metric_type: str, name: str, unit: Optional[str], attributes: Optional[dict]
 ) -> str:
@@ -153,6 +208,24 @@ def _series_key(
     return "\x00".join((metric_type, name, unit or "", attrs_part))
 
 
+def _build_exemplars(exemplar: Optional[dict]) -> dict:
+    # Spread into the data point, so a point without an exemplar keeps its
+    # exact pre-exemplar wire shape. OTLP/JSON encodes trace and span IDs as hex.
+    if exemplar is None:
+        return {}
+    return {
+        "exemplars": [
+            {
+                "filteredAttributes": [],
+                "timeUnixNano": _ms_to_unix_nano(exemplar["time_ms"]),
+                "asDouble": exemplar["value"],
+                "traceId": exemplar["trace_id"],
+                "spanId": exemplar["span_id"],
+            }
+        ]
+    }
+
+
 class _SeriesState:
     __slots__ = (
         "name",
@@ -163,6 +236,7 @@ class _SeriesState:
         "total",
         "last",
         "hist",
+        "exemplar",
     )
 
     def __init__(
@@ -195,6 +269,7 @@ def __init__(
         self.total: Optional[float] = None
         self.last: Optional[float] = None
         self.hist: Optional[dict] = None
+        self.exemplar: Optional[dict] = None
 
 
 class PostHogMetrics:
@@ -412,6 +487,7 @@ def _capture(
             # series identity too — otherwise two indistinguishable data points emit.
             attributes = {k: v for k, v in attributes.items() if v is not None}
         key = _series_key(metric_type, name, unit, attributes)
+        trace_context = _active_trace_context(self._client)
 
         with self._lock:
             self._reset_after_fork_locked()
@@ -446,6 +522,13 @@ def _capture(
                 )
 
             self._fold(state, float(value))
+            if trace_context is not None:
+                state.exemplar = {
+                    "trace_id": trace_context[0],
+                    "span_id": trace_context[1],
+                    "value": float(value),
+                    "time_ms": int(time.time() * 1000),
+                }
             self._arm_flush_timer()
 
     def _reinit_after_fork(self) -> None:
@@ -641,6 +724,9 @@ def _merge_window_back(self, window: dict) -> None:
                 self._series[key] = old
                 continue
             current.window_start_ms = min(current.window_start_ms, old.window_start_ms)
+            # The live window's exemplar is newer; the old one only fills a gap.
+            if current.exemplar is None:
+                current.exemplar = old.exemplar
             if current.type == "count":
                 current.total = (current.total or 0.0) + (old.total or 0.0)
             elif current.type == "histogram" and old.hist:
@@ -718,6 +804,7 @@ def _build_metrics(self, window: dict) -> list:
 
             attributes = _to_otlp_key_value_list(state.attributes or {})
             start_nano = _ms_to_unix_nano(state.window_start_ms)
+            exemplars = _build_exemplars(state.exemplar)
 
             if state.type == "count":
                 metric["sum"]["dataPoints"].append(
@@ -726,6 +813,7 @@ def _build_metrics(self, window: dict) -> list:
                         "startTimeUnixNano": start_nano,
                         "timeUnixNano": now_nano,
                         "asDouble": state.total or 0.0,
+                        **exemplars,
                     }
                 )
             elif state.type == "gauge":
@@ -734,6 +822,7 @@ def _build_metrics(self, window: dict) -> list:
                         "attributes": attributes,
                         "timeUnixNano": now_nano,
                         "asDouble": state.last or 0.0,
+                        **exemplars,
                     }
                 )
             elif state.hist is not None:
@@ -751,6 +840,7 @@ def _build_metrics(self, window: dict) -> list:
                         "max": state.hist["max"],
                         "bucketCounts": state.hist["bucket_counts"],
                         "explicitBounds": DEFAULT_HISTOGRAM_BOUNDS,
+                        **exemplars,
                     }
                 )
 

posthog/test/test_metrics.py

@@ -4,6 +4,8 @@
 from unittest import mock
 
 import pytest
+from opentelemetry import trace
+from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
 
 import posthog
 from posthog.client import Client
@@ -25,6 +27,21 @@ def client():
     c.metrics.reset()
 
 
+@pytest.fixture
+def traced_client():
+    # Tracing on, so start_span returns a recording span with its own span ID
+    # instead of a pass-through that echoes the parent's traceparent.
+    with mock.patch(
+        "posthog.tracing._transport._get_session", return_value=mock_session()
+    ):
+        c = Client(
+            FAKE_API_KEY, host="https://us.example.com", sync_mode=True, traces={}
+        )
+        yield c
+        c.metrics.reset()
+        c.shutdown()
+
+
 def mock_session(status_code=200):
     session = mock.Mock()
     session.post.return_value = mock.Mock(status_code=status_code)
@@ -586,6 +603,120 @@ def test_any_value_encoding_matches_js(self, client, value, expected):
         assert attr["value"] == expected
 
 
+OTEL_TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736"
+OTEL_SPAN_ID = "00f067aa0ba902b7"
+
+
+def otel_span(sampled):
+    return NonRecordingSpan(
+        SpanContext(
+            trace_id=int(OTEL_TRACE_ID, 16),
+            span_id=int(OTEL_SPAN_ID, 16),
+            is_remote=False,
+            trace_flags=TraceFlags(
+                TraceFlags.SAMPLED if sampled else TraceFlags.DEFAULT
+            ),
+        )
+    )
+
+
+def data_points_by_type(payload):
+    points = {}
+    for metric in metrics_from(payload):
+        for kind in ("sum", "gauge", "histogram"):
+            if kind in metric:
+                points[kind] = metric[kind]["dataPoints"][0]
+    return points
+
+
+def record_one_of_each(client):
+    client.metrics.count("c", 1)
+    client.metrics.count("c", 4)
+    client.metrics.gauge("g", 7)
+    client.metrics.histogram("h", 12)
+
+
+class TestMetricsExemplars:
+    def test_sampled_otel_span_stamps_last_sample_on_every_type(self, client):
+        with trace.use_span(otel_span(sampled=True)):
+            record_one_of_each(client)
+
+        payload, _, _ = flush_and_capture(client)
+        points = data_points_by_type(payload)
+
+        for kind, value in (("sum", 4), ("gauge", 7), ("histogram", 12)):
+            (exemplar,) = points[kind]["exemplars"]
+            assert exemplar["traceId"] == OTEL_TRACE_ID
+            assert exemplar["spanId"] == OTEL_SPAN_ID
+            assert exemplar["asDouble"] == value
+            assert exemplar["timeUnixNano"].isdigit()
+
+    def test_active_posthog_span_wins_over_outer_otel_span(self, traced_client):
+        parent_trace_id, parent_span_id = "a" * 32, "b" * 16
+        with trace.use_span(otel_span(sampled=True)):
+            with traced_client.start_span(
+                "job", parent=f"00-{parent_trace_id}-{parent_span_id}-01"
+            ) as span:
+                traced_client.metrics.count("c", 1)
+                _, trace_id, span_id, _ = span.traceparent().split("-")
+
+        payload, _, _ = flush_and_capture(traced_client)
+        (exemplar,) = data_points_by_type(payload)["sum"]["exemplars"]
+        # The recording span's own ID, in the inbound trace, not the parent's.
+        assert span_id not in (parent_span_id, OTEL_SPAN_ID)
+        assert (exemplar["traceId"], exemplar["spanId"]) == (parent_trace_id, span_id)
+
+    def test_innermost_posthog_span_is_the_exemplar(self, traced_client):
+        with traced_client.start_span("request"):
+            with traced_client.start_span("db.query") as child:
+                traced_client.metrics.histogram("h", 3)
+                _, trace_id, span_id, _ = child.traceparent().split("-")
+
+        payload, _, _ = flush_and_capture(traced_client)
+        (exemplar,) = data_points_by_type(payload)["histogram"]["exemplars"]
+        assert (exemplar["traceId"], exemplar["spanId"]) == (trace_id, span_id)
+
+    @pytest.mark.parametrize("context", ["no_span", "unsampled_otel_span"])
+    def test_no_exemplar_without_a_sampled_span(self, client, context):
+        if context == "no_span":
+            record_one_of_each(client)
+        else:
+            with trace.use_span(otel_span(sampled=False)):
+                record_one_of_each(client)
+
+        payload, _, _ = flush_and_capture(client)
+        for point in data_points_by_type(payload).values():
+            assert "exemplars" not in point
+
+    def test_unsampled_posthog_span_hides_a_sampled_outer_otel_span(
+        self, traced_client
+    ):
+        with trace.use_span(otel_span(sampled=True)):
+            with traced_client.start_span(
+                "job", parent=f"00-{'a' * 32}-{'b' * 16}-00"
+            ) as span:
+                assert span.traceparent().endswith("-00")
+                record_one_of_each(traced_client)
+
+        payload, _, _ = flush_and_capture(traced_client)
+        for point in data_points_by_type(payload).values():
+            assert "exemplars" not in point
+
+    def test_failed_flush_keeps_exemplar_for_the_retry(self, client):
+        with trace.use_span(otel_span(sampled=True)):
+            client.metrics.count("c", 1)
+        with mock.patch(
+            "posthog.metrics_capture._get_session", return_value=mock_session(503)
+        ):
+            client.metrics.flush()
+        client.metrics.count("c", 2)
+
+        payload, _, _ = flush_and_capture(client)
+        point = data_points_by_type(payload)["sum"]
+        assert point["asDouble"] == 3
+        assert point["exemplars"][0]["traceId"] == OTEL_TRACE_ID
+
+
 class TestMetricsFailureHandling:
     def test_fork_drops_inherited_window_and_rearms(self, client):
         # A forked child inherits a dead timer thread and the parent's window; without the