feat(perf): Add ability to put measurements directly on spans. (#2967)

getsentry · Apr 30, 2024 · 5130590 · 5130590
1 parent 842df5e
commit 5130590
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 35 deletions.
diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py
@@ -28,6 +28,45 @@
     # "critical" is an alias of "fatal" recognized by Relay
     LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
 
+    DurationUnit = Literal[
+        "nanosecond",
+        "microsecond",
+        "millisecond",
+        "second",
+        "minute",
+        "hour",
+        "day",
+        "week",
+    ]
+
+    InformationUnit = Literal[
+        "bit",
+        "byte",
+        "kilobyte",
+        "kibibyte",
+        "megabyte",
+        "mebibyte",
+        "gigabyte",
+        "gibibyte",
+        "terabyte",
+        "tebibyte",
+        "petabyte",
+        "pebibyte",
+        "exabyte",
+        "exbibyte",
+    ]
+
+    FractionUnit = Literal["ratio", "percent"]
+    MeasurementUnit = Union[DurationUnit, InformationUnit, FractionUnit, str]
+
+    MeasurementValue = TypedDict(
+        "MeasurementValue",
+        {
+            "value": float,
+            "unit": Optional[MeasurementUnit],
+        },
+    )
+
     Event = TypedDict(
         "Event",
         {
@@ -49,7 +88,7 @@
             "level": LogLevelStr,
             "logentry": Mapping[str, object],
             "logger": str,
-            "measurements": dict[str, object],
+            "measurements": dict[str, MeasurementValue],
             "message": str,
             "modules": dict[str, str],
             "monitor_config": Mapping[str, object],
@@ -118,37 +157,6 @@
     ]
     SessionStatus = Literal["ok", "exited", "crashed", "abnormal"]
 
-    DurationUnit = Literal[
-        "nanosecond",
-        "microsecond",
-        "millisecond",
-        "second",
-        "minute",
-        "hour",
-        "day",
-        "week",
-    ]
-
-    InformationUnit = Literal[
-        "bit",
-        "byte",
-        "kilobyte",
-        "kibibyte",
-        "megabyte",
-        "mebibyte",
-        "gigabyte",
-        "gibibyte",
-        "terabyte",
-        "tebibyte",
-        "petabyte",
-        "pebibyte",
-        "exabyte",
-        "exbibyte",
-    ]
-
-    FractionUnit = Literal["ratio", "percent"]
-    MeasurementUnit = Union[DurationUnit, InformationUnit, FractionUnit, str]
-
     ProfilerMode = Literal["sleep", "thread", "gevent", "unknown"]
 
     # Type of the metric.

diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -12,7 +12,6 @@
 )
 from sentry_sdk._types import TYPE_CHECKING
 
-
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping, MutableMapping
     from typing import Any
@@ -32,7 +31,12 @@
     R = TypeVar("R")
 
     import sentry_sdk.profiler
-    from sentry_sdk._types import Event, MeasurementUnit, SamplingContext
+    from sentry_sdk._types import (
+        Event,
+        MeasurementUnit,
+        SamplingContext,
+        MeasurementValue,
+    )
 
     class SpanKwargs(TypedDict, total=False):
         trace_id: str
@@ -189,6 +193,7 @@ class Span:
         "sampled",
         "op",
         "description",
+        "_measurements",
         "start_timestamp",
         "_start_timestamp_monotonic_ns",
         "status",
@@ -229,6 +234,7 @@ def __init__(
         self.status = status
         self.hub = hub
         self.scope = scope
+        self._measurements = {}  # type: Dict[str, MeasurementValue]
         self._tags = {}  # type: MutableMapping[str, str]
         self._data = {}  # type: Dict[str, Any]
         self._containing_transaction = containing_transaction
@@ -488,6 +494,10 @@ def set_status(self, value):
         # type: (str) -> None
         self.status = value
 
+    def set_measurement(self, name, value, unit=""):
+        # type: (str, float, MeasurementUnit) -> None
+        self._measurements[name] = {"value": value, "unit": unit}
+
     def set_thread(self, thread_id, thread_name):
         # type: (Optional[int], Optional[str]) -> None
 
@@ -598,6 +608,9 @@ def to_json(self):
             if metrics_summary:
                 rv["_metrics_summary"] = metrics_summary
 
+        if len(self._measurements) > 0:
+            rv["measurements"] = self._measurements
+
         tags = self._tags
         if tags:
             rv["tags"] = tags
@@ -674,7 +687,7 @@ def __init__(
         self.source = source
         self.sample_rate = None  # type: Optional[float]
         self.parent_sampled = parent_sampled
-        self._measurements = {}  # type: Dict[str, Any]
+        self._measurements = {}  # type: Dict[str, MeasurementValue]
         self._contexts = {}  # type: Dict[str, Any]
         self._profile = None  # type: Optional[sentry_sdk.profiler.Profile]
         self._baggage = baggage