gendesign/tradein-mvp/backend/app/services/proxy_pool.py

"""Пул прокси: подбор (lease), освобождение, health-трекинг (#2162).

АДДИТИВНО. Этот модуль реализует pick/lease/release/health-механику поверх таблицы
scrape_proxies (миграция 157, #2161). Ни один боевой скрейпер здесь НЕ подключается —
интеграция pick-из-пула вместо env-прокси это отдельные шаги P3/P4. Пока модуль
используется только health-checker'ом (run_proxy_healthcheck), который просто гоняет
ipify-пробу через каждый прокси и обновляет health-поля.

Семантика lease:
  - scrape_proxies.leased_by IS NULL  → прокси свободен.
  - leased_by = <run_id>              → занят run'ом.
  - leased_by = NON_RUN_LEASE_MARKER  → занят не-run вызовом (health-checker и т.п.),
                                        когда run_id не применим.
  acquire() берёт строку через FOR UPDATE SKIP LOCKED (конкурентные acquire не дерутся
  за одну строку — второй параллельный вызов пропустит залоченную и возьмёт следующую).

Health:
  - mark_health(ok=True)  → consecutive_fails=0, last_ok_at/last_check_at, exit_ip, latency.
  - mark_health(ok=False) → consecutive_fails += 1; при достижении DISABLE_THRESHOLD прокси
                            авто-disable (enabled=false), чтобы битый узел выпал из пула.
  - acquire отфильтровывает enabled=false И consecutive_fails >= MAX_FAILS.

psycopg v3 / SQLAlchemy text(): все параметры через CAST(:x AS type), НЕ :x::type.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass

import httpx
from sqlalchemy import text
from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)

__all__ = [
    "DISABLE_THRESHOLD",
    "MAX_CONSECUTIVE_FAILS",
    "NON_RUN_LEASE_MARKER",
    "STALE_LEASE_MINUTES",
    "ProxyLease",
    "acquire",
    "mark_health",
    "reap_stale_leases",
    "release",
    "run_proxy_healthcheck",
]

# ── Пороги ───────────────────────────────────────────────────────────────────
# Прокси с >= MAX_CONSECUTIVE_FAILS подряд-фейлами не выдаётся acquire'ом (даже если
# ещё enabled) — «карантин» до первого успешного health-check'а (mark_health сбросит
# счётчик в 0). Мягче, чем disable: узел может ожить.
MAX_CONSECUTIVE_FAILS = 3

# При достижении этого порога подряд-фейлов прокси авто-disable (enabled=false) —
# оператор включит вручную после разбора. Строго >= MAX_CONSECUTIVE_FAILS.
DISABLE_THRESHOLD = 5

# Lease старше этого времени считается протухшим (упавший sweep не вызвал release) и
# освобождается reap_stale_leases — иначе прокси навсегда «занят» мёртвым run'ом.
STALE_LEASE_MINUTES = 30

# Маркер lease для не-run вызовов (leased_by NOT NULL = занят, но это не id из scrape_runs).
NON_RUN_LEASE_MARKER = -1

# URL для health-пробы: возвращает exit-IP JSON'ом. Тот же эндпоинт, что и admin
# /scraper/health (_probe_current_ip).
_HEALTH_PROBE_URL = "https://api.ipify.org"
_HEALTH_PROBE_TIMEOUT_S = 10.0


@dataclass
class ProxyLease:
    """Арендованный прокси. url несёт схему (http:// / socks5://) — готов для httpx proxy=."""

    id: int
    url: str
    kind: str
    rotate_url: str | None


def acquire(db: Session, provider: str, *, run_id: int | None = None) -> ProxyLease | None:
    """Взять свободный здоровый прокси под провайдера (avito/cian/yandex/generic/any).

    SELECT ... FOR UPDATE SKIP LOCKED LIMIT 1 отбирает enabled-прокси с приемлемым
    health (consecutive_fails < MAX_CONSECUTIVE_FAILS), affinity=provider ИЛИ 'any',
    ещё не арендованный (leased_by IS NULL), предпочитая давно не проверенные
    (last_ok_at NULLS LAST). Затем помечает строку leased_by=run_id (или
    NON_RUN_LEASE_MARKER если run_id не задан) и коммитит.

    Конкурентные acquire не дерутся за одну строку: SKIP LOCKED пропускает залоченную
    другим вызовом строку, второй параллельный acquire берёт следующую свободную.

    Returns ProxyLease или None если свободных здоровых прокси нет.
    """
    lease_marker = run_id if run_id is not None else NON_RUN_LEASE_MARKER

    row = (
        db.execute(
            text(
                """
                SELECT id, url, kind, rotate_url
                FROM scrape_proxies
                WHERE enabled
                  AND consecutive_fails < CAST(:max_fails AS integer)
                  AND provider_affinity IN (:provider, 'any')
                  AND leased_by IS NULL
                ORDER BY last_ok_at NULLS LAST, id
                FOR UPDATE SKIP LOCKED
                LIMIT 1
                """
            ),
            {"max_fails": MAX_CONSECUTIVE_FAILS, "provider": provider},
        )
        .mappings()
        .fetchone()
    )
    if row is None:
        db.rollback()  # снять FOR UPDATE-транзакцию (ничего не залочено, но чисто)
        return None

    proxy_id = int(row["id"])
    db.execute(
        text(
            """
            UPDATE scrape_proxies
            SET leased_by = CAST(:run_id AS bigint), leased_at = now()
            WHERE id = CAST(:id AS bigint)
            """
        ),
        {"run_id": lease_marker, "id": proxy_id},
    )
    db.commit()
    logger.info(
        "proxy_pool: leased proxy id=%d provider=%s by=%s", proxy_id, provider, lease_marker
    )
    return ProxyLease(
        id=proxy_id,
        url=str(row["url"]),
        kind=str(row["kind"]),
        rotate_url=row["rotate_url"],
    )


def release(db: Session, proxy_id: int) -> None:
    """Освободить прокси (leased_by/leased_at → NULL). Идемпотентно (0-row если уже свободен)."""
    db.execute(
        text(
            """
            UPDATE scrape_proxies
            SET leased_by = NULL, leased_at = NULL
            WHERE id = CAST(:id AS bigint)
            """
        ),
        {"id": proxy_id},
    )
    db.commit()
    logger.info("proxy_pool: released proxy id=%d", proxy_id)


def mark_health(
    db: Session,
    proxy_id: int,
    ok: bool,
    *,
    exit_ip: str | None = None,
    latency_ms: int | None = None,
) -> None:
    """Записать результат health-check'а прокси.

    ok=True  → consecutive_fails обнуляется, обновляются last_ok_at/last_check_at/
               exit_ip/latency_ms.
    ok=False → consecutive_fails += 1; при достижении DISABLE_THRESHOLD прокси
               авто-disable (enabled=false). last_check_at обновляется в любом случае.
    """
    if ok:
        db.execute(
            text(
                """
                UPDATE scrape_proxies
                SET consecutive_fails = 0,
                    last_ok_at    = now(),
                    last_check_at = now(),
                    exit_ip       = CAST(:exit_ip AS text),
                    latency_ms    = CAST(:latency_ms AS integer),
                    updated_at    = now()
                WHERE id = CAST(:id AS bigint)
                """
            ),
            {"exit_ip": exit_ip, "latency_ms": latency_ms, "id": proxy_id},
        )
    else:
        # consecutive_fails+1 >= порог → enabled=false (авто-вывод битого узла).
        db.execute(
            text(
                """
                UPDATE scrape_proxies
                SET consecutive_fails = consecutive_fails + 1,
                    last_check_at     = now(),
                    enabled = CASE
                        WHEN consecutive_fails + 1 >= CAST(:disable_threshold AS integer)
                        THEN false ELSE enabled
                    END,
                    updated_at = now()
                WHERE id = CAST(:id AS bigint)
                """
            ),
            {"disable_threshold": DISABLE_THRESHOLD, "id": proxy_id},
        )
    db.commit()
    logger.info("proxy_pool: mark_health id=%d ok=%s exit_ip=%s", proxy_id, ok, exit_ip)


def reap_stale_leases(db: Session, older_than_minutes: int = STALE_LEASE_MINUTES) -> int:
    """Освободить lease'ы старше older_than_minutes (упавший sweep не вызвал release).

    Без этого прокси навсегда «занят» мёртвым run'ом и выпадает из пула. Returns число
    освобождённых прокси.
    """
    rows = db.execute(
        text(
            """
            UPDATE scrape_proxies
            SET leased_by = NULL, leased_at = NULL
            WHERE leased_by IS NOT NULL
              AND leased_at < now() - make_interval(mins => CAST(:mins AS integer))
            RETURNING id
            """
        ),
        {"mins": older_than_minutes},
    ).fetchall()
    db.commit()
    if rows:
        logger.warning("proxy_pool: reaped %d stale lease(s): %s", len(rows), [r.id for r in rows])
    return len(rows)


async def _probe_proxy(url: str) -> tuple[bool, str | None, int | None]:
    """GET ipify через прокси (timeout _HEALTH_PROBE_TIMEOUT_S).

    Returns (ok, exit_ip, latency_ms). ok=False + (None, None) при любой ошибке.
    url несёт схему (http:// / socks5://) — httpx[socks] обрабатывает оба.
    """
    started = time.monotonic()
    try:
        async with httpx.AsyncClient(proxy=url, timeout=_HEALTH_PROBE_TIMEOUT_S) as client:
            resp = await client.get(_HEALTH_PROBE_URL, params={"format": "json"})
            resp.raise_for_status()
            ip = resp.json().get("ip")
        latency_ms = int((time.monotonic() - started) * 1000)
        return True, (str(ip) if ip else None), latency_ms
    except Exception:
        logger.warning("proxy_pool: health probe failed proxy=%s", _mask(url), exc_info=True)
        return False, None, None


def _mask(url: str) -> str:
    """Скрыть пароль в proxy-url для логов (scheme://user:***@host)."""
    if "@" not in url or "//" not in url:
        return url
    scheme, rest = url.split("//", 1)
    creds, host = rest.split("@", 1)
    if ":" in creds:
        user, _pwd = creds.split(":", 1)
        creds = f"{user}:***"
    return f"{scheme}//{creds}@{host}"


async def run_proxy_healthcheck(db: Session) -> dict[str, int]:
    """Периодический health-check всех enabled-прокси пула (#2162).

    Сначала reap_stale_leases (освобождает протухшие lease'ы), затем для каждого
    enabled-прокси гоняет ipify-пробу через сам прокси и пишет результат через
    mark_health (успех → сброс fails + свежий exit_ip/latency; фейл → инкремент,
    авто-disable при DISABLE_THRESHOLD).

    Пробы идут последовательно — пул небольшой (десятки узлов), а параллельный залп на
    один и тот же upstream-endpoint (ipify) не нужен. Returns counters
    {reaped, checked, ok, failed}.
    """
    reaped = reap_stale_leases(db)

    proxies = (
        db.execute(
            text(
                """
                SELECT id, url, kind
                FROM scrape_proxies
                WHERE enabled
                ORDER BY id
                """
            )
        )
        .mappings()
        .all()
    )

    checked = 0
    ok_count = 0
    failed = 0
    for row in proxies:
        proxy_id = int(row["id"])
        url = str(row["url"])
        ok, exit_ip, latency_ms = await _probe_proxy(url)
        mark_health(db, proxy_id, ok, exit_ip=exit_ip, latency_ms=latency_ms)
        checked += 1
        if ok:
            ok_count += 1
        else:
            failed += 1

    logger.info(
        "proxy_pool: healthcheck done — reaped=%d checked=%d ok=%d failed=%d",
        reaped,
        checked,
        ok_count,
        failed,
    )
    return {"reaped": reaped, "checked": checked, "ok": ok_count, "failed": failed}