anti-plagiarism/services/api/app/core/rate_limiter.py

"""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