167 lines
4.9 KiB
Python
167 lines
4.9 KiB
Python
import os
|
|
import random
|
|
import threading
|
|
from datetime import datetime, timezone
|
|
from typing import Optional, List, Callable, TYPE_CHECKING, Any, Union
|
|
|
|
from sentry_sdk.utils import format_timestamp, safe_repr
|
|
from sentry_sdk.envelope import Envelope, Item, PayloadRef
|
|
|
|
if TYPE_CHECKING:
|
|
from sentry_sdk._types import Metric
|
|
|
|
|
|
class MetricsBatcher:
|
|
MAX_METRICS_BEFORE_FLUSH = 1000
|
|
MAX_METRICS_BEFORE_DROP = 10_000
|
|
FLUSH_WAIT_TIME = 5.0
|
|
|
|
def __init__(
|
|
self,
|
|
capture_func, # type: Callable[[Envelope], None]
|
|
record_lost_func, # type: Callable[..., None]
|
|
):
|
|
# type: (...) -> None
|
|
self._metric_buffer = [] # type: List[Metric]
|
|
self._capture_func = capture_func
|
|
self._record_lost_func = record_lost_func
|
|
self._running = True
|
|
self._lock = threading.Lock()
|
|
|
|
self._flush_event = threading.Event() # type: threading.Event
|
|
|
|
self._flusher = None # type: Optional[threading.Thread]
|
|
self._flusher_pid = None # type: Optional[int]
|
|
|
|
def _ensure_thread(self):
|
|
# type: (...) -> bool
|
|
if not self._running:
|
|
return False
|
|
|
|
pid = os.getpid()
|
|
if self._flusher_pid == pid:
|
|
return True
|
|
|
|
with self._lock:
|
|
if self._flusher_pid == pid:
|
|
return True
|
|
|
|
self._flusher_pid = pid
|
|
|
|
self._flusher = threading.Thread(target=self._flush_loop)
|
|
self._flusher.daemon = True
|
|
|
|
try:
|
|
self._flusher.start()
|
|
except RuntimeError:
|
|
self._running = False
|
|
return False
|
|
|
|
return True
|
|
|
|
def _flush_loop(self):
|
|
# type: (...) -> None
|
|
while self._running:
|
|
self._flush_event.wait(self.FLUSH_WAIT_TIME + random.random())
|
|
self._flush_event.clear()
|
|
self._flush()
|
|
|
|
def add(
|
|
self,
|
|
metric, # type: Metric
|
|
):
|
|
# type: (...) -> None
|
|
if not self._ensure_thread() or self._flusher is None:
|
|
return None
|
|
|
|
with self._lock:
|
|
if len(self._metric_buffer) >= self.MAX_METRICS_BEFORE_DROP:
|
|
self._record_lost_func(
|
|
reason="queue_overflow",
|
|
data_category="trace_metric",
|
|
quantity=1,
|
|
)
|
|
return None
|
|
|
|
self._metric_buffer.append(metric)
|
|
if len(self._metric_buffer) >= self.MAX_METRICS_BEFORE_FLUSH:
|
|
self._flush_event.set()
|
|
|
|
def kill(self):
|
|
# type: (...) -> None
|
|
if self._flusher is None:
|
|
return
|
|
|
|
self._running = False
|
|
self._flush_event.set()
|
|
self._flusher = None
|
|
|
|
def flush(self):
|
|
# type: (...) -> None
|
|
self._flush()
|
|
|
|
@staticmethod
|
|
def _metric_to_transport_format(metric):
|
|
# type: (Metric) -> Any
|
|
def format_attribute(val):
|
|
# type: (Union[int, float, str, bool]) -> Any
|
|
if isinstance(val, bool):
|
|
return {"value": val, "type": "boolean"}
|
|
if isinstance(val, int):
|
|
return {"value": val, "type": "integer"}
|
|
if isinstance(val, float):
|
|
return {"value": val, "type": "double"}
|
|
if isinstance(val, str):
|
|
return {"value": val, "type": "string"}
|
|
return {"value": safe_repr(val), "type": "string"}
|
|
|
|
res = {
|
|
"timestamp": metric["timestamp"],
|
|
"trace_id": metric["trace_id"],
|
|
"name": metric["name"],
|
|
"type": metric["type"],
|
|
"value": metric["value"],
|
|
"attributes": {
|
|
k: format_attribute(v) for (k, v) in metric["attributes"].items()
|
|
},
|
|
}
|
|
|
|
if metric.get("span_id") is not None:
|
|
res["span_id"] = metric["span_id"]
|
|
|
|
if metric.get("unit") is not None:
|
|
res["unit"] = metric["unit"]
|
|
|
|
return res
|
|
|
|
def _flush(self):
|
|
# type: (...) -> Optional[Envelope]
|
|
|
|
envelope = Envelope(
|
|
headers={"sent_at": format_timestamp(datetime.now(timezone.utc))}
|
|
)
|
|
with self._lock:
|
|
if len(self._metric_buffer) == 0:
|
|
return None
|
|
|
|
envelope.add_item(
|
|
Item(
|
|
type="trace_metric",
|
|
content_type="application/vnd.sentry.items.trace-metric+json",
|
|
headers={
|
|
"item_count": len(self._metric_buffer),
|
|
},
|
|
payload=PayloadRef(
|
|
json={
|
|
"items": [
|
|
self._metric_to_transport_format(metric)
|
|
for metric in self._metric_buffer
|
|
]
|
|
}
|
|
),
|
|
)
|
|
)
|
|
self._metric_buffer.clear()
|
|
|
|
self._capture_func(envelope)
|
|
return envelope
|