RandomZ / app /rate_limit /redis_limiter.py
StormShadow308's picture
feat: async pipeline, job queue, generation hardening, and docs
732b14f
Raw
History Blame Contribute Delete
1.27 kB
"""Redis ZSET sliding-window rate limiter (shared across API replicas)."""
from __future__ import annotations
import time
import uuid
from redis.asyncio import Redis
class RedisSlidingWindowLimiter:
"""Per-tenant sliding window using a Redis sorted set."""
def __init__(
self,
redis: Redis,
*,
key_prefix: str,
max_requests: int,
window_secs: float,
) -> None:
self._redis = redis
self._prefix = key_prefix
self._max = max_requests
self._window = window_secs
async def is_allowed(self, tenant_key: str) -> bool:
now = time.time()
cutoff = now - self._window
rk = f"{self._prefix}:{tenant_key}"
member = f"{now}:{uuid.uuid4().hex}"
pipe = self._redis.pipeline()
pipe.zremrangebyscore(rk, 0, cutoff)
pipe.zadd(rk, {member: now})
pipe.zcard(rk)
pipe.expire(rk, int(self._window) + 1)
results = await pipe.execute()
count = int(results[2])
if count > self._max:
await self._redis.zrem(rk, member)
return False
return True
async def reset(self, tenant_key: str) -> None:
await self._redis.delete(f"{self._prefix}:{tenant_key}")