68 lines
1.4 KiB
Python
68 lines
1.4 KiB
Python
import threading
|
|
from typing import Any, Protocol
|
|
|
|
import httpx2
|
|
|
|
|
|
class ApiRuntimeStats:
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._total_received = 0
|
|
|
|
def record_received(self) -> None:
|
|
with self._lock:
|
|
self._total_received += 1
|
|
|
|
def snapshot(self) -> dict[str, int]:
|
|
with self._lock:
|
|
return {
|
|
"total_received": self._total_received,
|
|
}
|
|
|
|
|
|
class WorkerStatsUnavailable(Exception):
|
|
pass
|
|
|
|
|
|
class WorkerStatsProvider(Protocol):
|
|
async def get_stats(
|
|
self,
|
|
) -> dict[str, Any]:
|
|
...
|
|
|
|
async def close(self) -> None:
|
|
...
|
|
|
|
|
|
class HttpWorkerStatsClient:
|
|
def __init__(
|
|
self,
|
|
url: str,
|
|
timeout_seconds: float,
|
|
) -> None:
|
|
self._url = url
|
|
|
|
self._client = httpx2.AsyncClient(
|
|
timeout=timeout_seconds,
|
|
)
|
|
|
|
async def get_stats(
|
|
self,
|
|
) -> dict[str, Any]:
|
|
try:
|
|
response = await self._client.get(
|
|
self._url
|
|
)
|
|
|
|
response.raise_for_status()
|
|
|
|
return response.json()
|
|
|
|
except httpx2.HTTPError as exc:
|
|
raise WorkerStatsUnavailable(
|
|
f"Unable to retrieve worker stats: {exc}"
|
|
) from exc
|
|
|
|
async def close(self) -> None:
|
|
await self._client.aclose()
|