1"""Redis helpers: API cache (prefix `aia:api:`), distributed locks (`aia:lock:`), heartbeats."""2from __future__ import annotations34import asyncio5import logging6from collections.abc import AsyncIterator7from contextlib import asynccontextmanager8from typing import Any910import orjson11from redis.asyncio import Redis1213from aiatlas.config import settings1415log = logging.getLogger(__name__)16_redis: Redis | None = None171819def redis() -> Redis:20 global _redis21 if _redis is None:22 _redis = Redis.from_url(settings.redis_url, decode_responses=False, socket_connect_timeout=3, socket_timeout=5)23 return _redis242526async def cache_get(key: str) -> Any | None:27 try:28 raw = await redis().get(f"aia:api:{key}")29 return orjson.loads(raw) if raw else None30 except Exception: # noqa: BLE00131 return None323334async def cache_set(key: str, value: Any, ttl_s: int = 300) -> None:35 try:36 await redis().set(f"aia:api:{key}", orjson.dumps(value, default=str), ex=ttl_s)37 except Exception: # noqa: BLE00138 pass394041async def cache_invalidate(prefix: str = "") -> int:42 n = 043 try:44 async for key in redis().scan_iter(match=f"aia:api:{prefix}*", count=500):45 await redis().delete(key)46 n += 147 except Exception: # noqa: BLE00148 pass49 return n505152@asynccontextmanager53async def lock(name: str, ttl_s: int = 3600) -> AsyncIterator[bool]:54 """Best-effort distributed lock. Yields False when someone else holds it (callers must skip)."""55 key = f"aia:lock:{name}"56 acquired = False57 try:58 acquired = bool(await redis().set(key, b"1", nx=True, ex=ttl_s))59 except Exception as exc: # noqa: BLE00160 log.warning("redis unavailable, proceeding without lock", extra={"error": str(exc)})61 acquired = True62 try:63 yield acquired64 finally:65 if acquired:66 try:67 await redis().delete(key)68 except Exception: # noqa: BLE00169 pass707172async def heartbeat(service: str, payload: dict[str, Any], ttl_s: int = 180) -> None:73 try:74 await redis().set(f"aia:heartbeat:{service}", orjson.dumps(payload, default=str), ex=ttl_s)75 except Exception: # noqa: BLE00176 pass777879async def heartbeats() -> dict[str, Any]:80 out: dict[str, Any] = {}81 try:82 async for key in redis().scan_iter(match="aia:heartbeat:*"):83 raw = await redis().get(key)84 if raw:85 out[key.decode().rsplit(":", 1)[-1]] = orjson.loads(raw)86 except Exception: # noqa: BLE00187 pass88 return out899091async def close() -> None:92 global _redis93 if _redis is not None:94 await _redis.aclose()95 _redis = None96 await asyncio.sleep(0)979899__all__ = ["redis", "cache_get", "cache_set", "cache_invalidate", "lock", "heartbeat", "heartbeats", "close"]100