ForcePilot/backend/package/yuxi/channel/message/metrics.py

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)