WIP autoscaler agent

agent/nix_builder_autoscaler/metrics.py

@@ -0,0 +1,103 @@
+"""In-memory Prometheus metrics registry.
+
+No prometheus_client dependency — formats text manually.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+
+def _labels_key(labels: dict[str, str]) -> tuple[tuple[str, str], ...]:
+    return tuple(sorted(labels.items()))
+
+
+def _format_labels(labels: dict[str, str]) -> str:
+    if not labels:
+        return ""
+    parts = ",".join(f'{k}="{v}"' for k, v in sorted(labels.items()))
+    return "{" + parts + "}"
+
+
+class MetricsRegistry:
+    """Thread-safe in-memory metrics store with Prometheus text output."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._gauges: dict[str, dict[tuple[tuple[str, str], ...], float]] = {}
+        self._counters: dict[str, dict[tuple[tuple[str, str], ...], float]] = {}
+        self._histograms: dict[str, dict[tuple[tuple[str, str], ...], Any]] = {}
+
+    def gauge(self, name: str, labels: dict[str, str], value: float) -> None:
+        """Set a gauge value."""
+        key = _labels_key(labels)
+        with self._lock:
+            if name not in self._gauges:
+                self._gauges[name] = {}
+            self._gauges[name][key] = value
+
+    def counter(self, name: str, labels: dict[str, str], increment: float = 1.0) -> None:
+        """Increment a counter."""
+        key = _labels_key(labels)
+        with self._lock:
+            if name not in self._counters:
+                self._counters[name] = {}
+            self._counters[name][key] = self._counters[name].get(key, 0.0) + increment
+
+    def histogram_observe(self, name: str, labels: dict[str, str], value: float) -> None:
+        """Record a histogram observation.
+
+        Uses fixed buckets: 0.01, 0.05, 0.1, 0.5, 1, 5, 10, 30, 60, 120, +Inf.
+        """
+        key = _labels_key(labels)
+        buckets = (0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 120.0)
+        with self._lock:
+            if name not in self._histograms:
+                self._histograms[name] = {}
+            if key not in self._histograms[name]:
+                self._histograms[name][key] = {
+                    "labels": labels,
+                    "buckets": {b: 0 for b in buckets},
+                    "sum": 0.0,
+                    "count": 0,
+                }
+            entry = self._histograms[name][key]
+            entry["sum"] += value
+            entry["count"] += 1
+            for b in buckets:
+                if value <= b:
+                    entry["buckets"][b] += 1
+
+    def render(self) -> str:
+        """Render all metrics in Prometheus text exposition format."""
+        lines: list[str] = []
+        with self._lock:
+            for name, series in sorted(self._gauges.items()):
+                lines.append(f"# TYPE {name} gauge")
+                for key, val in sorted(series.items()):
+                    labels = dict(key)
+                    lines.append(f"{name}{_format_labels(labels)} {val}")
+
+            for name, series in sorted(self._counters.items()):
+                lines.append(f"# TYPE {name} counter")
+                for key, val in sorted(series.items()):
+                    labels = dict(key)
+                    lines.append(f"{name}{_format_labels(labels)} {val}")
+
+            for name, series in sorted(self._histograms.items()):
+                lines.append(f"# TYPE {name} histogram")
+                for _key, entry in sorted(series.items()):
+                    labels = entry["labels"]
+                    cumulative = 0
+                    for b, count in sorted(entry["buckets"].items()):
+                        cumulative += count
+                        le_labels = {**labels, "le": str(b)}
+                        lines.append(f"{name}_bucket{_format_labels(le_labels)} {cumulative}")
+                    inf_labels = {**labels, "le": "+Inf"}
+                    lines.append(f"{name}_bucket{_format_labels(inf_labels)} {entry['count']}")
+                    lines.append(f"{name}_sum{_format_labels(labels)} {entry['sum']}")
+                    lines.append(f"{name}_count{_format_labels(labels)} {entry['count']}")
+
+        lines.append("")
+        return "\n".join(lines)