"""Redis-based rate limiter с атомарными Lua-скриптами. Проблема наивного подхода (GET → проверка → INCR): - Race condition: 10 конкурентных запросов могут одновременно пройти GET, увидеть значение ниже лимита и все инкрементировать. Решение: один Lua-скрипт выполняется атомарно на стороне Redis. Redis гарантирует, что между командами внутри скрипта нет других операций. """ import logging from datetime import datetime, timezone from app.core.redis_client import get_redis logger = logging.getLogger(__name__) # ─── Лимиты по тарифам ──────────────────────────────────────────────────────── PLAN_LIMITS: dict[str, dict[str, int | None]] = { "free": { "search_per_day": 10, "summarize_per_month": 3, "plagiarism_per_month": 1, "concurrent": 1, }, "student": { "search_per_day": None, # None = безлимит "summarize_per_month": 30, "plagiarism_per_month": 10, "concurrent": 2, }, "premium": { "search_per_day": None, "summarize_per_month": None, "plagiarism_per_month": 50, "concurrent": 5, }, "science": { "search_per_day": None, "summarize_per_month": None, "plagiarism_per_month": None, "concurrent": 10, }, } ACTION_TO_LIMIT: dict[str, tuple[str, str]] = { "search": ("search_per_day", "day"), "summarize": ("summarize_per_month", "month"), "plagiarism":("plagiarism_per_month", "month"), # gost не ограничен — не в маппинге } # ─── Lua-скрипты ────────────────────────────────────────────────────────────── # Атомарная проверка + инкремент лимита. # Возвращает [current_value, allowed] где allowed = 1 если OK, 0 если превышен. _LUA_CHECK_AND_INCR = """ local key = KEYS[1] local limit = tonumber(ARGV[1]) local ttl = tonumber(ARGV[2]) local current = tonumber(redis.call('GET', key) or '0') if current >= limit then return {current, 0} end local new_val = redis.call('INCR', key) -- Устанавливаем TTL только при первом инкременте (когда ключ только что создан) if new_val == 1 then redis.call('EXPIRE', key, ttl) end return {new_val, 1} """ # Атомарная проверка + инкремент счётчика одновременных задач. # Возвращает 1 если слот получен, 0 если превышен лимит. _LUA_ACQUIRE_CONCURRENT = """ local key = KEYS[1] local limit = tonumber(ARGV[1]) local ttl = tonumber(ARGV[2]) local current = tonumber(redis.call('GET', key) or '0') if current >= limit then return 0 end redis.call('INCR', key) redis.call('EXPIRE', key, ttl) return 1 """ def _period_suffix(period: str) -> str: now = datetime.now(timezone.utc) return now.strftime("%Y-%m-%d") if period == "day" else now.strftime("%Y-%m") async def check_and_increment_limit( user_id: int, action: str, plan: str, ) -> dict: """ Атомарно проверить лимит и инкрементировать счётчик. Returns: {allowed, current, limit, remaining, reset_at} """ limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) if action not in ACTION_TO_LIMIT: return {"allowed": True, "current": 0, "limit": None, "remaining": None} limit_key, period = ACTION_TO_LIMIT[action] limit_value = limits.get(limit_key) if limit_value is None: # безлимит return {"allowed": True, "current": 0, "limit": None, "remaining": None} period_str = _period_suffix(period) redis_key = f"rl:{user_id}:{action}:{period_str}" ttl = 86_400 if period == "day" else 86_400 * 32 r = get_redis() result = await r.eval(_LUA_CHECK_AND_INCR, 1, redis_key, limit_value, ttl) current, allowed = int(result[0]), bool(result[1]) return { "allowed": allowed, "current": current, "limit": limit_value, "remaining": max(0, limit_value - current) if allowed else 0, "reset_at": period_str, } async def acquire_concurrent_slot(user_id: int, plan: str) -> bool: """ Атомарно захватить слот одновременной задачи. Returns: True — слот получен (задачу можно создавать). False — все слоты заняты. """ limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) max_concurrent = limits.get("concurrent", 1) r = get_redis() result = await r.eval( _LUA_ACQUIRE_CONCURRENT, 1, f"concurrent:{user_id}", max_concurrent, 3_600, # TTL 1 час — автосброс если воркер упал не освободив слот ) return bool(result) async def release_concurrent_slot(user_id: int) -> None: """Освободить слот одновременной задачи после завершения.""" r = get_redis() key = f"concurrent:{user_id}" # DECR безопасен: Redis не уходит в отрицательные значения если мы контролируем acquire current = await r.get(key) if current and int(current) > 0: await r.decr(key)