import time from collections import defaultdict from collections.abc import Callable from dataclasses import dataclass, field from enum import StrEnum class MetricType(StrEnum): COUNTER = "counter" GAUGE = "gauge" HISTOGRAM = "histogram" @dataclass(slots=True) class _Counter: value: int = 0 def inc(self, amount: int = 1) -> None: self.value += amount @dataclass(slots=True) class _Gauge: value: float = 0.0 def set(self, value: float) -> None: self.value = value def inc(self, amount: float = 1.0) -> None: self.value += amount def dec(self, amount: float = 1.0) -> None: self.value -= amount @dataclass(slots=True) class _Histogram: buckets: list[float] values: list[int] = field(default_factory=list) _sum: float = 0.0 _count: int = 0 def observe(self, value: float) -> None: self._sum += value self._count += 1 while len(self.values) < len(self.buckets): self.values.append(0) for i, bound in enumerate(self.buckets): if value <= bound: self.values[i] += 1 return def quantile(self, q: float) -> float | None: if self._count == 0: return None if not self.values: return None target_rank = q * self._count cumulative = 0 for i, (bound, count) in enumerate(zip(self.buckets, self.values)): cumulative += count if cumulative >= target_rank: if i == 0: lower_bound = 0.0 else: lower_bound = self.buckets[i - 1] upper_bound = bound prev_cumulative = cumulative - count fraction = (target_rank - prev_cumulative) / max(count, 1) return lower_bound + fraction * (upper_bound - lower_bound) return float(self.buckets[-1]) if self.buckets else None class MetricsRegistry: def __init__(self) -> None: self._label_values: dict[str, dict[tuple[str, ...], object]] = defaultdict(dict) def counter(self, name: str, label_keys: tuple[str, ...] = ()) -> Callable[..., None]: def inc(labels: dict[str, str] | None = None, amount: int = 1) -> None: key = self._resolve_key(labels, label_keys) counter = self._label_values[name].get(key) if counter is None: counter = _Counter() self._label_values[name][key] = counter counter.inc(amount) return inc def gauge(self, name: str, label_keys: tuple[str, ...] = ()) -> Callable[..., None]: def set(value: float, labels: dict[str, str] | None = None) -> None: key = self._resolve_key(labels, label_keys) gauge = self._label_values[name].get(key) if gauge is None: gauge = _Gauge() self._label_values[name][key] = gauge gauge.set(value) return set def histogram(self, name: str, buckets: list[float], label_keys: tuple[str, ...] = ()) -> Callable[..., None]: def observe(value: float, labels: dict[str, str] | None = None) -> None: key = self._resolve_key(labels, label_keys) hist = self._label_values[name].get(key) if hist is None: hist = _Histogram(buckets=buckets) self._label_values[name][key] = hist hist.observe(value) return observe @staticmethod def _resolve_key(labels: dict[str, str] | None, label_keys: tuple[str, ...]) -> tuple[str, ...]: if not labels: return tuple("" for _ in label_keys) return tuple(labels.get(k, "") for k in label_keys) def snapshot(self) -> dict: result: dict = {} for metric_name, entries in self._label_values.items(): result[metric_name] = {} for label_tuple, metric in entries.items(): if isinstance(metric, _Counter): result[metric_name][str(label_tuple)] = {"type": "counter", "value": metric.value} elif isinstance(metric, _Gauge): result[metric_name][str(label_tuple)] = {"type": "gauge", "value": metric.value} elif isinstance(metric, _Histogram): result[metric_name][str(label_tuple)] = { "type": "histogram", "buckets": metric.buckets, "values": metric.values, "sum": metric._sum, "count": metric._count, "p50": metric.quantile(0.50), "p95": metric.quantile(0.95), "p99": metric.quantile(0.99), } return result registry = MetricsRegistry() channel_messages_total = registry.counter( "channel_messages_total", ("channel_type", "status"), ) channel_dispatch_duration_ms = registry.histogram( "channel_dispatch_duration_ms", buckets=[5, 25, 50, 100, 250, 500, 1000, 2500, 5000], label_keys=("channel_type",), ) channel_agent_duration_ms = registry.histogram( "channel_agent_duration_ms", buckets=[100, 500, 1000, 2500, 5000, 10000, 30000, 60000], label_keys=("channel_type",), ) channel_rate_limit_rejects_total = registry.counter( "channel_rate_limit_rejects_total", ("channel_type",), ) channel_messages_inflight = registry.gauge( "channel_messages_inflight", ) def record_message(channel_type: str, status: str) -> None: channel_messages_total(labels={"channel_type": channel_type, "status": status}) def record_dispatch_duration_ms(channel_type: str, duration_ms: float) -> None: channel_dispatch_duration_ms(duration_ms, labels={"channel_type": channel_type}) def record_agent_duration_ms(channel_type: str, duration_ms: float) -> None: channel_agent_duration_ms(duration_ms, labels={"channel_type": channel_type}) def record_rate_limit_reject(channel_type: str) -> None: channel_rate_limit_rejects_total(labels={"channel_type": channel_type}) def set_inflight(count: int) -> None: channel_messages_inflight(count) class MetricsTimer: def __init__(self, on_finish: Callable[[float], None]) -> None: self._on_finish = on_finish self._start = 0.0 def __enter__(self) -> "MetricsTimer": self._start = time.monotonic() return self def __exit__(self, *args) -> None: elapsed = (time.monotonic() - self._start) * 1000 self._on_finish(elapsed) async def __aenter__(self) -> "MetricsTimer": self._start = time.monotonic() return self async def __aexit__(self, *args) -> None: elapsed = (time.monotonic() - self._start) * 1000 self._on_finish(elapsed)