initial commit

bandwidth-worker/worker.py

@@ -0,0 +1,145 @@
+"""bandwidth-worker: polls frps dashboard, writes deltas to Postgres via PostgREST."""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+from datetime import datetime, timezone
+from typing import Any
+
+import httpx
+import redis
+
+POLL_INTERVAL = int(os.environ.get("POLL_INTERVAL_SECONDS", "60"))
+FRPS_API = os.environ.get("FRPS_API_URL", "http://frps:7500/api/proxy/http")
+FRPS_USER = os.environ["FRPS_DASHBOARD_USER"]
+FRPS_PASS = os.environ["FRPS_DASHBOARD_PASS"]
+SUPABASE_URL = os.environ.get("SUPABASE_URL", "http://supabase-kong:8000")
+SERVICE_ROLE = os.environ["SUPABASE_SERVICE_ROLE_KEY"]
+REDIS_URL = os.environ["REDIS_URL"]
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(message)s",
+    stream=sys.stdout,
+)
+log = logging.getLogger("bandwidth-worker")
+
+rds = redis.from_url(REDIS_URL, decode_responses=True)
+last_seen: dict[str, int] = {}
+
+
+def _headers(extra: dict[str, str] | None = None) -> dict[str, str]:
+    h = {
+        "apikey": SERVICE_ROLE,
+        "authorization": f"Bearer {SERVICE_ROLE}",
+        "content-type": "application/json",
+    }
+    if extra:
+        h.update(extra)
+    return h
+
+
+def fetch_proxies(client: httpx.Client) -> list[dict[str, Any]]:
+    resp = client.get(FRPS_API, auth=(FRPS_USER, FRPS_PASS), timeout=10.0)
+    resp.raise_for_status()
+    return resp.json().get("proxies") or []
+
+
+def previous_total(subdomain: str) -> int:
+    if subdomain in last_seen:
+        return last_seen[subdomain]
+    try:
+        v = rds.get(f"tunnel:bytes_total:{subdomain}")
+        if v is not None:
+            return int(v)
+    except Exception as e:
+        log.warning("redis get failed for %s: %s", subdomain, e)
+    return 0
+
+
+def persist_total(subdomain: str, total: int) -> None:
+    last_seen[subdomain] = total
+    try:
+        rds.set(f"tunnel:bytes_total:{subdomain}", total)
+    except Exception as e:
+        log.warning("redis set failed for %s: %s", subdomain, e)
+
+
+def record_delta(client: httpx.Client, subdomain: str, delta: int) -> None:
+    r = client.post(
+        f"{SUPABASE_URL}/rest/v1/usage_samples",
+        headers=_headers({"prefer": "return=minimal"}),
+        json={"subdomain": subdomain, "bytes_delta": delta},
+        timeout=10.0,
+    )
+    if r.status_code >= 300:
+        log.error("usage_samples insert failed %s: %s", r.status_code, r.text)
+        return
+
+    cur = client.get(
+        f"{SUPABASE_URL}/rest/v1/tunnels?select=bytes_used&subdomain=eq.{subdomain}",
+        headers=_headers(),
+        timeout=10.0,
+    )
+    if cur.status_code >= 300:
+        log.error("tunnels select failed %s: %s", cur.status_code, cur.text)
+        return
+    rows = cur.json()
+    if not rows:
+        log.warning("tunnel row missing for subdomain=%s; sample retained", subdomain)
+        return
+    current = int(rows[0].get("bytes_used") or 0)
+    upd = client.patch(
+        f"{SUPABASE_URL}/rest/v1/tunnels?subdomain=eq.{subdomain}",
+        headers=_headers({"prefer": "return=minimal"}),
+        json={
+            "bytes_used": current + delta,
+            "last_seen_at": datetime.now(timezone.utc).isoformat(),
+        },
+        timeout=10.0,
+    )
+    if upd.status_code >= 300:
+        log.error("tunnels update failed %s: %s", upd.status_code, upd.text)
+
+
+def poll_once(client: httpx.Client) -> None:
+    proxies = fetch_proxies(client)
+    log.info("poll: %d proxies", len(proxies))
+    for p in proxies:
+        name = p.get("name") or ""
+        if not name:
+            continue
+        traffic_in = int(p.get("today_traffic_in") or 0)
+        traffic_out = int(p.get("today_traffic_out") or 0)
+        total = traffic_in + traffic_out
+        prev = previous_total(name)
+        delta = total - prev
+        if delta < 0:
+            delta = total  # daily counter reset
+        if delta == 0:
+            persist_total(name, total)
+            continue
+        try:
+            record_delta(client, name, delta)
+            persist_total(name, total)
+            log.info("recorded delta=%d for %s (total=%d)", delta, name, total)
+        except Exception as e:
+            log.error("record_delta failed for %s: %s", name, e)
+
+
+def main() -> None:
+    log.info("bandwidth-worker starting; interval=%ss api=%s", POLL_INTERVAL, FRPS_API)
+    with httpx.Client() as client:
+        while True:
+            try:
+                poll_once(client)
+            except Exception as e:
+                log.error("poll failed: %s", e)
+            time.sleep(POLL_INTERVAL)
+
+
+if __name__ == "__main__":
+    main()