feat: troppa roba davvero

backend/app/services/mitmproxy/health.py

@@ -1,8 +1,11 @@
 import asyncio
 import ssl
+import time
 from dataclasses import dataclass
+from datetime import UTC, datetime
 
-from models.proxy import HealthCheckDocument, ProxyDocument
+from app.models.health import HealthCheckEntry, HealthSnapshot
+from app.models.proxy import HealthCheckDocument, ProxyDocument
 
 
 @dataclass(slots=True)
@@ -13,33 +16,48 @@ class Target:
     protocol: str
 
 
-async def check_proxy(proxy: ProxyDocument) -> dict[str, object]:
+async def check_proxy(proxy: ProxyDocument) -> HealthSnapshot:
+    timestamp = datetime.now(UTC)
     if not proxy.healthcheck.enabled:
-        return {"enabled": False, "overall": "disabled", "checks": []}
+        return HealthSnapshot(
+            proxy_port=proxy.proxy_port,
+            timestamp=timestamp,
+            enabled=False,
+            overall="disabled",
+            checks=[],
+        )
 
     upstream_port = proxy.service.upstream_port or proxy.service.port
-    checks = await asyncio.gather(
-        check_target(
-            Target(
-                name="upstream",
-                host=proxy.service.host,
-                port=upstream_port,
-                protocol=health_protocol(proxy),
+    checks = list(
+        await asyncio.gather(
+            check_target(
+                Target(
+                    name="upstream",
+                    host=proxy.service.host,
+                    port=upstream_port,
+                    protocol=health_protocol(proxy),
+                ),
+                proxy.healthcheck,
             ),
-            proxy.healthcheck,
-        ),
-        check_target(
-            Target(
-                name="proxy",
-                host=listen_host_for_check(proxy.listen_host),
-                port=proxy.proxy_port,
-                protocol=health_protocol(proxy),
+            check_target(
+                Target(
+                    name="proxy",
+                    host=listen_host_for_check(proxy.listen_host),
+                    port=proxy.proxy_port,
+                    protocol=health_protocol(proxy),
+                ),
+                proxy.healthcheck,
             ),
-            proxy.healthcheck,
-        ),
+        )
+    )
+    overall = "healthy" if all(check.healthy for check in checks) else "unhealthy"
+    return HealthSnapshot(
+        proxy_port=proxy.proxy_port,
+        timestamp=timestamp,
+        enabled=True,
+        overall=overall,
+        checks=checks,
     )
-    overall = "healthy" if all(check["healthy"] for check in checks) else "unhealthy"
-    return {"enabled": True, "overall": overall, "checks": checks}
 
 
 def health_protocol(proxy: ProxyDocument) -> str:
@@ -55,7 +73,7 @@ def listen_host_for_check(host: str) -> str:
 async def check_target(
     target: Target,
     config: HealthCheckDocument,
-) -> dict[str, object]:
+) -> HealthCheckEntry:
     try:
         match target.protocol:
             case "http":
@@ -81,7 +99,8 @@ async def check_http(
     config: HealthCheckDocument,
     *,
     use_tls: bool,
-) -> dict[str, object]:
+) -> HealthCheckEntry:
+    started = time.perf_counter()
     reader, writer = await open_connection(target, config, use_tls=use_tls)
     try:
         path = config.path if config.path.startswith("/") else f"/{config.path}"
@@ -101,7 +120,13 @@ async def check_http(
     status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else None
     expected = config.expected_status
     healthy = status == expected if expected is not None else status is not None
-    return result(target, healthy, None, {"status": status})
+    return result(
+        target,
+        healthy,
+        None,
+        {"status": status},
+        benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
+    )
 
 
 async def check_tcp(
@@ -109,14 +134,20 @@ async def check_tcp(
     config: HealthCheckDocument,
     *,
     use_tls: bool,
-) -> dict[str, object]:
+) -> HealthCheckEntry:
+    started = time.perf_counter()
     _reader, writer = await open_connection(target, config, use_tls=use_tls)
     writer.close()
     await writer.wait_closed()
-    return result(target, True)
+    return result(
+        target,
+        True,
+        benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
+    )
 
 
-async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, object]:
+async def check_udp(target: Target, config: HealthCheckDocument) -> HealthCheckEntry:
+    started = time.perf_counter()
     loop = asyncio.get_running_loop()
     transport, protocol = await loop.create_datagram_endpoint(
         lambda: DatagramProbe(config.expect_response),
@@ -128,7 +159,13 @@ async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, ob
         if not config.expect_response:
             return result(target, True, None, {"mode": "send-only"})
         data = await asyncio.wait_for(protocol.response, config.timeout_seconds)
-        return result(target, True, None, {"bytes": len(data)})
+        return result(
+            target,
+            True,
+            None,
+            {"bytes": len(data)},
+            benchmark_ms=elapsed_ms(started) if config.benchmark_enabled else None,
+        )
     finally:
         transport.close()
 
@@ -159,7 +196,7 @@ class DatagramProbe(asyncio.DatagramProtocol):
         if not expect_response:
             self.response.set_result(b"")
 
-    def datagram_received(self, data, _addr):
+    def datagram_received(self, data, addr):
         if not self.response.done():
             self.response.set_result(data)
 
@@ -173,13 +210,19 @@ def result(
     healthy: bool,
     error: str | None = None,
     details: dict[str, object] | None = None,
-) -> dict[str, object]:
-    return {
-        "target": target.name,
-        "host": target.host,
-        "port": target.port,
-        "protocol": target.protocol,
-        "healthy": healthy,
-        "error": error,
-        "details": details or {},
-    }
+    benchmark_ms: float | None = None,
+) -> HealthCheckEntry:
+    return HealthCheckEntry(
+        target=target.name,
+        host=target.host,
+        port=target.port,
+        protocol=target.protocol,
+        healthy=healthy,
+        error=error,
+        duration_ms=benchmark_ms,
+        details=details or {},
+    )
+
+
+def elapsed_ms(started: float) -> float:
+    return round((time.perf_counter() - started) * 1000, 2)