"""Rate limiter: дневные/месячные лимиты — в Redis, лимит одновременных задач — в БД. Дневные/месячные лимиты (Redis, атомарные Lua-скрипты): Проблема наивного подхода (GET → проверка → INCR): - Race condition: 10 конкурентных запросов могут одновременно пройти GET, увидеть значение ниже лимита и все инкрементировать. Решение: один Lua-скрипт выполняется атомарно на стороне Redis. Redis гарантирует, что между командами внутри скрипта нет других операций. Лимит одновременных задач — НЕ Redis-счётчик. Раньше был acquire/release-счётчик (INCR при создании задачи, DECR при завершении) — но release_concurrent_slot() никогда не вызывался ни из одного воркера, так что счётчик только рос и лимит превышался навсегда (до часового TTL-автосброса), даже когда все задачи юзера давно завершены. Вместо ручного счётчика, который может рассинхронизироваться с реальностью, считаем активные задачи прямо по Task.status в Postgres — рассинхронизации тогда не может быть в принципе. """ import logging from datetime import UTC, datetime from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession 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} """ def _period_suffix(period: str) -> str: now = datetime.now(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 check_concurrent_limit(db: AsyncSession, user_id: int, plan: str) -> bool: """ Проверить лимит одновременных задач по факту в БД (queued/processing). Возможен редкий race (два запроса одновременно оба видят N-1 активных и оба проходят) — на практике не критично для этого лимита (защита от злоупотребления, не от превышения на единицу), а взамен исключён класс багов "счётчик разошёлся с реальностью и завис навсегда". Returns: True — лимит не превышен, задачу можно создавать. """ from app.models.task import Task # избегаем circular import на уровне модуля limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) max_concurrent = limits.get("concurrent", 1) result = await db.execute( select(func.count()) .select_from(Task) .where(Task.user_id == user_id, Task.status.in_(("queued", "processing"))) ) current = result.scalar_one() return current < max_concurrent