118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
import asyncio
|
|
import logging
|
|
|
|
from app.core.config import settings
|
|
from app.models.health import HealthSnapshot
|
|
from app.models.proxy import ProxyDocument
|
|
from app.repositories.health import record_health_sample
|
|
from app.repositories.proxy import list_proxies
|
|
from app.services.mitmproxy.health import check_proxy
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class HealthMonitor:
|
|
def __init__(self):
|
|
self.task: asyncio.Task[None] | None = None
|
|
self.subscribers: dict[int, set[asyncio.Queue[HealthSnapshot]]] = {}
|
|
|
|
async def start(self) -> None:
|
|
if self.task is None or self.task.done():
|
|
self.task = asyncio.create_task(self.run())
|
|
|
|
async def stop(self) -> None:
|
|
if self.task is None:
|
|
return
|
|
self.task.cancel()
|
|
await asyncio.gather(self.task, return_exceptions=True)
|
|
self.task = None
|
|
|
|
async def run(self) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
next_cycle = loop.time()
|
|
while True:
|
|
delay = next_cycle - loop.time()
|
|
if delay > 0:
|
|
await asyncio.sleep(delay)
|
|
|
|
cycle_started = loop.time()
|
|
try:
|
|
await self.run_cycle()
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
logger.exception("Health monitor cycle failed")
|
|
next_cycle = loop.time() + min(
|
|
1.0,
|
|
float(settings.HEALTHCHECK_INTERVAL_SECONDS),
|
|
)
|
|
continue
|
|
|
|
next_cycle = cycle_started + settings.HEALTHCHECK_INTERVAL_SECONDS
|
|
|
|
async def run_cycle(self) -> None:
|
|
proxies = [
|
|
proxy
|
|
for proxy in await list_proxies()
|
|
if proxy.status == "running" and proxy.healthcheck.enabled
|
|
]
|
|
async with asyncio.TaskGroup() as group:
|
|
for proxy in proxies:
|
|
group.create_task(self.run_check(proxy))
|
|
|
|
async def run_check(self, proxy: ProxyDocument) -> None:
|
|
try:
|
|
snapshot = await check_proxy(proxy)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
logger.exception(
|
|
"Health check failed for proxy %s",
|
|
proxy.proxy_port,
|
|
)
|
|
return
|
|
|
|
try:
|
|
await record_health_sample(snapshot)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
logger.exception(
|
|
"Health snapshot persistence failed for proxy %s",
|
|
proxy.proxy_port,
|
|
)
|
|
|
|
await self.publish(proxy.proxy_port, snapshot)
|
|
|
|
async def publish(self, proxy_port: int, snapshot: HealthSnapshot) -> None:
|
|
for queue in list(self.subscribers.get(proxy_port, set())):
|
|
try:
|
|
queue.put_nowait(snapshot)
|
|
except asyncio.QueueFull:
|
|
logger.debug(
|
|
"Dropped health update for slow subscriber on proxy %s",
|
|
proxy_port,
|
|
)
|
|
continue
|
|
|
|
def subscribe(self, proxy_port: int) -> asyncio.Queue[HealthSnapshot]:
|
|
queue: asyncio.Queue[HealthSnapshot] = asyncio.Queue(maxsize=8)
|
|
listeners = self.subscribers.setdefault(proxy_port, set())
|
|
listeners.add(queue)
|
|
return queue
|
|
|
|
def unsubscribe(
|
|
self,
|
|
proxy_port: int,
|
|
queue: asyncio.Queue[HealthSnapshot],
|
|
) -> None:
|
|
listeners = self.subscribers.get(proxy_port)
|
|
if listeners is None:
|
|
return
|
|
listeners.discard(queue)
|
|
if not listeners:
|
|
self.subscribers.pop(proxy_port, None)
|
|
|
|
|
|
health_monitor = HealthMonitor()
|