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

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