from __future__ import annotations import asyncio import time from contextlib import asynccontextmanager from uuid import uuid4 from app.config import settings class ProviderLimitTimeout(TimeoutError): pass @asynccontextmanager async def provider_limit(name: str, limit: int | None = None, wait_timeout: float | None = None, ttl: int | None = None): """Best-effort Redis distributed semaphore. If Redis is disabled, it becomes a no-op. This keeps local development simple. """ limit = limit or 0 wait_timeout = wait_timeout if wait_timeout is not None else settings.PROVIDER_LIMIT_WAIT_TIMEOUT_SECONDS ttl = ttl or settings.PROVIDER_LIMIT_TOKEN_TTL_SECONDS if limit <= 0 or not settings.REDIS_URL: yield return import redis.asyncio as redis client = redis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True) key = f"provider_limit:{name}" token = str(uuid4()) acquired = False deadline = time.monotonic() + wait_timeout try: while time.monotonic() < deadline: now = time.time() pipe = client.pipeline() pipe.zremrangebyscore(key, 0, now - ttl) pipe.zcard(key) _, count = await pipe.execute() if count < limit: added = await client.zadd(key, {token: now}, nx=True) await client.expire(key, ttl) if added: acquired = True break await asyncio.sleep(0.2) if not acquired: raise ProviderLimitTimeout(f"provider limit exceeded: {name}") yield finally: if acquired: try: await client.zrem(key, token) except Exception: pass await client.aclose()