Add up counter/histogram metrics in spans (#1099)

Author: alexmojaki

logfire/_internal/config.py

@@ -235,14 +235,14 @@ class PydanticPlugin:
 
 @dataclass
 class MetricsOptions:
-    """Configuration of metrics.
-
-    This only has one option for now, but it's a place to add more related options in the future.
-    """
+    """Configuration of metrics."""
 
     additional_readers: Sequence[MetricReader] = ()
     """Sequence of metric readers to be used in addition to the default which exports metrics to Logfire's API."""
 
+    collect_in_spans: bool = False
+    """Experimental setting to add up the values of counter and histogram metrics in active spans."""
+
 
 @dataclass
 class CodeSource:
@@ -1097,7 +1097,7 @@ def exit_open_spans():  # pragma: no cover
                 # The reason that spans may be lingering open is that they're in suspended generator frames.
                 # Apart from here, they will be ended when the generator is garbage collected
                 # as the interpreter shuts down, but that's too late.
-                for span in list(OPEN_SPANS):
+                for span in list(OPEN_SPANS.values()):
                     # TODO maybe we should be recording something about what happened here?
                     span.end()
                     # Interpreter shutdown may trigger another call to .end(),

logfire/_internal/metrics.py

@@ -21,8 +21,12 @@
     UpDownCounter,
 )
 from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider
+from opentelemetry.trace import get_current_span
 from opentelemetry.util.types import Attributes
 
+from .tracer import _LogfireWrappedSpan  # type: ignore
+from .utils import handle_internal_errors
+
 try:
     # This only exists in opentelemetry-sdk>=1.23.0
     from opentelemetry.metrics import _Gauge
@@ -247,6 +251,12 @@ def on_meter_set(self, meter: Meter) -> None:
     def _create_real_instrument(self, meter: Meter) -> InstrumentT:
         """Create an instance of the real instrument. Implement this."""
 
+    @handle_internal_errors
+    def _increment_span_metric(self, amount: float, attributes: Attributes | None = None):
+        span = get_current_span()
+        if isinstance(span, _LogfireWrappedSpan):
+            span.increment_metric(self._name, attributes or {}, amount)
+
 
 class _ProxyAsynchronousInstrument(_ProxyInstrument[InstrumentT], ABC):
     def __init__(
@@ -271,6 +281,7 @@ def add(
         *args: Any,
         **kwargs: Any,
     ) -> None:
+        self._increment_span_metric(amount, attributes)
         self._instrument.add(amount, attributes, *args, **kwargs)
 
     def _create_real_instrument(self, meter: Meter) -> Counter:
@@ -285,6 +296,7 @@ def record(
         *args: Any,
         **kwargs: Any,
     ) -> None:
+        self._increment_span_metric(amount, attributes)
         self._instrument.record(amount, attributes, *args, **kwargs)
 
     def _create_real_instrument(self, meter: Meter) -> Histogram:

logfire/_internal/tracer.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import json
 import sys
 import traceback
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Callable, cast
-from weakref import WeakKeyDictionary, WeakSet
+from weakref import WeakKeyDictionary, WeakValueDictionary
 
 import opentelemetry.trace as trace_api
 from opentelemetry import context as context_api
@@ -43,7 +45,7 @@
     ValidationError = None
 
 
-OPEN_SPANS: WeakSet[_LogfireWrappedSpan] = WeakSet()
+OPEN_SPANS: WeakValueDictionary[tuple[int, int], _LogfireWrappedSpan] = WeakValueDictionary()
 
 
 @dataclass
@@ -112,6 +114,23 @@ def force_flush(self, timeout_millis: int = 30000) -> bool:
             return True  # pragma: no cover
 
 
+@dataclass
+class SpanMetric:
+    details: dict[tuple[tuple[str, otel_types.AttributeValue], ...], float] = field(
+        default_factory=lambda: defaultdict(int)
+    )
+
+    def dump(self):
+        return {
+            'details': [{'attributes': dict(attributes), 'total': total} for attributes, total in self.details.items()],
+            'total': sum(total for total in self.details.values()),
+        }
+
+    def increment(self, attributes: Mapping[str, otel_types.AttributeValue], value: float):
+        key = tuple(sorted(attributes.items()))
+        self.details[key] += value
+
+
 @dataclass(eq=False)
 class _LogfireWrappedSpan(trace_api.Span, ReadableSpan):
     """A span that wraps another span and overrides some behaviors in a logfire-specific way.
@@ -124,14 +143,24 @@ class _LogfireWrappedSpan(trace_api.Span, ReadableSpan):
 
     span: Span
     ns_timestamp_generator: Callable[[], int]
+    record_metrics: bool
+    metrics: dict[str, SpanMetric] = field(default_factory=lambda: defaultdict(SpanMetric))
 
     def __post_init__(self):
-        OPEN_SPANS.add(self)
+        OPEN_SPANS[self._open_spans_key()] = self
 
     def end(self, end_time: int | None = None) -> None:
-        OPEN_SPANS.discard(self)
+        with handle_internal_errors:
+            OPEN_SPANS.pop(self._open_spans_key(), None)
+            if self.metrics:
+                self.span.set_attribute(
+                    'logfire.metrics', json.dumps({name: metric.dump() for name, metric in self.metrics.items()})
+                )
         self.span.end(end_time or self.ns_timestamp_generator())
 
+    def _open_spans_key(self):
+        return _open_spans_key(self.span.get_span_context())
+
     def get_span_context(self) -> SpanContext:
         return self.span.get_span_context()
 
@@ -175,6 +204,14 @@ def record_exception(
         timestamp = timestamp or self.ns_timestamp_generator()
         record_exception(self.span, exception, attributes=attributes, timestamp=timestamp, escaped=escaped)
 
+    def increment_metric(self, name: str, attributes: Mapping[str, otel_types.AttributeValue], value: float) -> None:
+        if not self.is_recording() or not self.record_metrics:
+            return
+
+        self.metrics[name].increment(attributes, value)
+        if self.parent and (parent := OPEN_SPANS.get(_open_spans_key(self.parent))):
+            parent.increment_metric(name, attributes, value)
+
     def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None:
         if self.is_recording():
             if isinstance(exc_value, BaseException):
@@ -187,6 +224,10 @@ def __getattr__(self, name: str) -> Any:
             return getattr(self.span, name)
 
 
+def _open_spans_key(ctx: SpanContext) -> tuple[int, int]:
+    return ctx.trace_id, ctx.span_id
+
+
 @dataclass
 class _ProxyTracer(Tracer):
     """A tracer that wraps another internal tracer allowing it to be re-assigned."""
@@ -216,7 +257,11 @@ def start_span(
         record_exception: bool = True,
         set_status_on_exception: bool = True,
     ) -> Span:
-        start_time = start_time or self.provider.config.advanced.ns_timestamp_generator()
+        config = self.provider.config
+        ns_timestamp_generator = config.advanced.ns_timestamp_generator
+        record_metrics: bool = not isinstance(config.metrics, (bool, type(None))) and config.metrics.collect_in_spans
+
+        start_time = start_time or ns_timestamp_generator()
 
         # Make a copy of the attributes since this method can be called by arbitrary external code,
         # e.g. third party instrumentation.
@@ -241,7 +286,8 @@ def start_span(
             )
         return _LogfireWrappedSpan(
             span,
-            ns_timestamp_generator=self.provider.config.advanced.ns_timestamp_generator,
+            ns_timestamp_generator=ns_timestamp_generator,
+            record_metrics=record_metrics,
         )
 
     # This means that `with start_as_current_span(...):`

logfire/integrations/pydantic.py

@@ -127,16 +127,16 @@ def wrapped_validator(input_data: Any, *args: Any, **kwargs: Any) -> Any:
                 try:
                     result = validator(input_data, *args, **kwargs)
                 except ValidationError as error:
-                    self._count_validation(success=False)
                     self._on_error_span(span, error)
+                    self._count_validation(success=False)
                     raise
                 except Exception as exception:
-                    self._count_validation(success=False)
                     self._on_exception_span(span, exception)
+                    self._count_validation(success=False)
                     raise
                 else:
-                    self._count_validation(success=True)
                     self._on_success(span, result)
+                    self._count_validation(success=True)
                     return result
 
         elif self._record == 'failure':

tests/conftest.py

@@ -101,6 +101,7 @@ def config(config_kwargs: dict[str, Any], metrics_reader: InMemoryMetricReader)
         **config_kwargs,
         metrics=logfire.MetricsOptions(
             additional_readers=[metrics_reader],
+            collect_in_spans=True,
         ),
     )
     # sanity check: there are no active spans

tests/otel_integrations/test_asgi.py

@@ -2,6 +2,7 @@
 
 import contextlib
 
+from dirty_equals import IsFloat, IsInt
 from inline_snapshot import snapshot
 from opentelemetry.propagate import inject
 from starlette.applications import Starlette
@@ -34,7 +35,7 @@ def homepage(_: Request):
     assert response.status_code == 200
     assert response.text == 'middleware test'
 
-    assert exporter.exported_spans_as_dict() == snapshot(
+    assert exporter.exported_spans_as_dict(parse_json_attributes=True) == snapshot(
         [
             {
                 'name': 'inside request handler',
@@ -98,6 +99,70 @@ def homepage(_: Request):
                     'logfire.msg_template': 'outside request handler',
                     'logfire.msg': 'outside request handler',
                     'logfire.span_type': 'span',
+                    'logfire.metrics': {
+                        'http.server.duration': {
+                            'details': [
+                                {
+                                    'attributes': {
+                                        'http.flavor': '1.1',
+                                        'http.host': 'testserver',
+                                        'http.method': 'GET',
+                                        'http.scheme': 'http',
+                                        'http.server_name': 'testserver',
+                                        'http.status_code': 200,
+                                        'net.host.port': 80,
+                                    },
+                                    'total': IsInt(),
+                                }
+                            ],
+                            'total': IsInt(),
+                        },
+                        'http.server.request.duration': {
+                            'details': [
+                                {
+                                    'attributes': {
+                                        'http.request.method': 'GET',
+                                        'http.response.status_code': 200,
+                                        'network.protocol.version': '1.1',
+                                        'url.scheme': 'http',
+                                    },
+                                    'total': IsFloat(),
+                                }
+                            ],
+                            'total': IsFloat(),
+                        },
+                        'http.server.response.size': {
+                            'details': [
+                                {
+                                    'attributes': {
+                                        'http.flavor': '1.1',
+                                        'http.host': 'testserver',
+                                        'http.method': 'GET',
+                                        'http.scheme': 'http',
+                                        'http.server_name': 'testserver',
+                                        'http.status_code': 200,
+                                        'net.host.port': 80,
+                                    },
+                                    'total': 15,
+                                }
+                            ],
+                            'total': 15,
+                        },
+                        'http.server.response.body.size': {
+                            'details': [
+                                {
+                                    'attributes': {
+                                        'http.request.method': 'GET',
+                                        'http.response.status_code': 200,
+                                        'network.protocol.version': '1.1',
+                                        'url.scheme': 'http',
+                                    },
+                                    'total': 15,
+                                }
+                            ],
+                            'total': 15,
+                        },
+                    },
                 },
             },
         ]