60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
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()
|