diff --git a/app/api/init_api.py b/app/api/init_api.py index 91b7cd755..5c25dc4dd 100644 --- a/app/api/init_api.py +++ b/app/api/init_api.py @@ -25,6 +25,7 @@ import app.bg_loops import app.settings import app.state +import app.telemetry import app.utils from app.api import api_router # type: ignore[attr-defined] from app.api import domains @@ -83,6 +84,8 @@ async def lifespan(asgi_app: BanchoAPI) -> AsyncIterator[Never]: Ansi.LRED, ) + app.telemetry.hook_database_calls() + await app.state.services.database.connect() await app.state.services.redis.initialize() diff --git a/app/telemetry.py b/app/telemetry.py new file mode 100644 index 000000000..0973e16d1 --- /dev/null +++ b/app/telemetry.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import dataclasses +import hashlib +import json +import platform +import time +from dataclasses import dataclass +from typing import Any + +import psutil +from httpx import AsyncClient +from tenacity import retry +from tenacity.stop import stop_after_attempt + +from app import settings + +MIN_SLOW_QUERY_SECONDS = 5.0 + +HTTP_CLIENT = AsyncClient() + + +@dataclass +class SoftwareInfo: + version: str + domain: str + + +@dataclass +class SystemInfo: + system: str + node: str + release: str + version: str + machine: str + processor: str + + +@dataclass +class LanguageInfo: + python_version: str + build_no: str + build_date: str + python_compiler: str + python_implementation: str + + +@dataclass +class SystemLoadInfo: + cpu_1min_average: float + cpu_5min_average: float + cpu_15min_average: float + + +@dataclass +class TelemetryEventReport: + software_info: SoftwareInfo + system_info: SystemInfo + language_info: LanguageInfo + system_load_info: SystemLoadInfo + event_data: dict[str, Any] + + +@retry(reraise=True, stop=stop_after_attempt(3)) +async def report_event(event_data: dict[str, Any]) -> None: + cpu_1min_average, cpu_5min_average, cpu_15min_average = psutil.getloadavg() + event = TelemetryEventReport( + software_info=SoftwareInfo( + version=settings.VERSION, + domain=settings.DOMAIN, + ), + system_info=SystemInfo( + system=platform.system(), + node=platform.node(), + release=platform.release(), + version=platform.version(), + machine=platform.machine(), + processor=platform.processor(), + ), + language_info=LanguageInfo( + python_version=platform.python_version(), + build_no=platform.python_build()[0], + build_date=platform.python_build()[1], + python_compiler=platform.python_compiler(), + python_implementation=platform.python_implementation(), + ), + system_load_info=SystemLoadInfo( + cpu_1min_average=cpu_1min_average, + cpu_5min_average=cpu_5min_average, + cpu_15min_average=cpu_15min_average, + ), + event_data=event_data, + ) + request_data = dataclasses.asdict(event) + idempotency_key = hashlib.sha256( + json.dumps(request_data, sort_keys=True).encode("utf-8"), + ).hexdigest() + response = await HTTP_CLIENT.post( + url="https://telemetry.cmyui.xyz/report", + headers={"Idempotency-Key": idempotency_key}, + json=request_data, + ) + response.raise_for_status() + return None + + +def hook_database_calls() -> None: + def _wrap_database_call(func): + async def wrapper(*args, **kwargs): + start_time = time.perf_counter() + response = await func(*args, **kwargs) + end_time = time.perf_counter() + + seconds_elapsed = end_time - start_time + + if seconds_elapsed >= MIN_SLOW_QUERY_SECONDS: + event_data = { + "query": kwargs.get("query", args and args[0]), + "seconds_elapsed": seconds_elapsed, + } + await report_event(event_data) + + return response + + return wrapper + + import app.state.services + + for attr in ("execute", "execute_many", "fetch_one", "fetch_all"): + unwrapped_func = getattr(app.state.services.database, attr) + wrapped_func = _wrap_database_call(unwrapped_func) + setattr(app.state.services.database, attr, wrapped_func) + + return None