swarm-house/simulator/monitoring/exporter.py

"""Prometheus exporter over the simulator's Parquet output.

Periodically scans DATA_DIR with DuckDB and exposes fleet statistics as
/metrics. Scans only the most recent flight partitions by default so CPU
stays bounded as the lake grows across pod restarts.
"""

from __future__ import annotations

import os
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import duckdb

sys.path.insert(0, str(Path(__file__).resolve().parent))
from lake import flight_window, iter_parquet_files, parquet_reader  # noqa: E402

DATA_DIR = Path(os.environ.get("DATA_DIR", "./data"))
PORT = int(os.environ.get("EXPORTER_PORT", "9105"))
SCAN_INTERVAL_S = float(os.environ.get("SCAN_INTERVAL_S", "15"))

_lock = threading.Lock()
_payload = "# swarm exporter starting\n"
_last_scan_s = 0.0


def _q(con: duckdb.DuckDBPyConnection, sql: str) -> list[tuple]:
    try:
        return con.sql(sql).fetchall()
    except duckdb.Error:
        return []


def collect() -> str:
    global _last_scan_s
    started = time.monotonic()
    con = duckdb.connect()
    lines: list[str] = []

    def metric(name: str, help_text: str, mtype: str, rows: list[str]) -> None:
        lines.append(f"# HELP {name} {help_text}")
        lines.append(f"# TYPE {name} {mtype}")
        lines.extend(rows)

    for ds, name in (
        ("telemetry", "swarm_rows_total"),
        ("detections", "swarm_detections_total"),
        ("state", "swarm_state_frames_total"),
    ):
        reader = parquet_reader(DATA_DIR, ds)
        if ds == "telemetry":
            metric(
                name, f"Rows in dataset={ds} (recent {flight_window()} flights)", "gauge",
                [f'{name}{{drone="{d}",sensor="{s}"}} {n}'
                 for d, s, n in _q(con, f"SELECT drone, sensor, count(*) FROM read_parquet({reader}) GROUP BY 1,2")],
            )
        elif ds == "detections":
            metric(
                name, "Detection events per drone and class", "gauge",
                [f'{name}{{drone="{d}",cls="{c}"}} {n}'
                 for d, c, n in _q(con, f"SELECT drone, cls, count(*) FROM read_parquet({reader}) GROUP BY 1,2")],
            )
        else:
            metric(
                name, "Pose broadcast frames per drone and direction", "gauge",
                [f'{name}{{drone="{d}",direction="{dr}"}} {n}'
                 for d, dr, n in _q(con, f"SELECT drone, direction, count(*) FROM read_parquet({reader}) GROUP BY 1,2")],
            )

    telem = parquet_reader(DATA_DIR, "telemetry")
    metric(
        "swarm_battery_pct", "Latest battery level per drone", "gauge",
        [f'swarm_battery_pct{{drone="{d}"}} {v}'
         for d, v in _q(con, f"""
            SELECT drone, arg_max(level_pct, ts_ns)
            FROM read_parquet({telem})
            WHERE sensor='battery' GROUP BY drone""")],
    )
    metric(
        "swarm_rssi_dbm", "Latest RSSI per drone-peer link", "gauge",
        [f'swarm_rssi_dbm{{drone="{d}",peer="{p}"}} {v}'
         for d, p, v in _q(con, f"""
            SELECT drone, peer_id, arg_max(rssi_dbm, ts_ns)
            FROM read_parquet({telem})
            WHERE sensor='rssi' GROUP BY drone, peer_id""")],
    )

    files = iter_parquet_files(DATA_DIR)
    metric(
        "swarm_parquet_bytes", "Bytes on disk per dataset (recent flights)", "gauge",
        [f'swarm_parquet_bytes{{dataset="{ds}"}} {sum(f.stat().st_size for f in files if f"dataset={ds}" in str(f))}'
         for ds in ("telemetry", "detections", "state")],
    )
    metric("swarm_parquet_files", "Parquet files scanned (recent flights)", "gauge",
           [f"swarm_parquet_files {len(files)}"])
    metric("swarm_metric_flight_window", "Flight partitions included per dataset", "gauge",
           [f"swarm_metric_flight_window {flight_window()}"])

    _last_scan_s = time.monotonic() - started
    metric("swarm_exporter_scan_duration_seconds", "Wall time of the last metrics scan", "gauge",
           [f"swarm_exporter_scan_duration_seconds {_last_scan_s:.4f}"])
    con.close()
    return "\n".join(lines) + "\n"


def scanner() -> None:
    global _payload
    while True:
        started = time.monotonic()
        try:
            payload = collect()
        except Exception as exc:
            payload = f"# collect error: {exc}\n"
        with _lock:
            _payload = payload
        time.sleep(max(1.0, SCAN_INTERVAL_S - (time.monotonic() - started)))


class Handler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:  # noqa: N802 — http.server API
        if self.path != "/metrics":
            self.send_response(404)
            self.end_headers()
            return
        with _lock:
            body = _payload.encode()
        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *_args: object) -> None:
        pass


if __name__ == "__main__":
    threading.Thread(target=scanner, daemon=True).start()
    print(
        f"swarm exporter on :{PORT}/metrics, scanning last {flight_window()} flights "
        f"every {SCAN_INTERVAL_S}s under {DATA_DIR}"
    )
    ThreadingHTTPServer(("", PORT), Handler).serve_forever()