eth_hedge_sim/backend/app/live/rate_limit.py

"""实盘交易限流：私有 REST 冷却 + 失败退避。"""

from __future__ import annotations

import logging
import threading
import time
from typing import Any

logger = logging.getLogger(__name__)

_DEFAULT_429_SEC = 20.0
_DEFAULT_418_SEC = 120.0
_INTERVAL_MIN = 0.2
_INTERVAL_MAX = 30.0


def resolve_live_order_interval_sec() -> float:
    """读取前端可配的 LIVE 下单最小间隔（秒），默认 1。"""
    try:
        from ..config import get_settings
        from ..models.db import get_db

        s = get_settings()
        default = float(s.live_order_interval_sec)
        raw = get_db().get_setting("live_order_interval_sec", str(default))
        v = float(raw if raw not in (None, "") else default)
        if v != v:  # NaN
            return 1.0
        return max(_INTERVAL_MIN, min(_INTERVAL_MAX, v))
    except Exception:
        return 1.0


class RateLimitError(RuntimeError):
    """处于限流/冷却中，调用方应退避，勿立即重试下单。"""

    def __init__(self, message: str, *, retry_after: float = 0.0) -> None:
        super().__init__(message)
        self.retry_after = float(retry_after)


class TradeThrottle:
    """按通道节流：最小间隔 + 418/429 冷却。"""

    def __init__(
        self,
        name: str,
        *,
        min_interval_sec: float = 1.0,
        cooldown_429_sec: float = _DEFAULT_429_SEC,
        cooldown_418_sec: float = _DEFAULT_418_SEC,
    ) -> None:
        self.name = name
        self.min_interval_sec = float(min_interval_sec)
        self.cooldown_429_sec = float(cooldown_429_sec)
        self.cooldown_418_sec = float(cooldown_418_sec)
        self._lock = threading.Lock()
        self._last_at = 0.0
        self._cool_until = 0.0

    def remaining_cooldown(self) -> float:
        with self._lock:
            return max(0.0, self._cool_until - time.monotonic())

    def before_request(self) -> None:
        """请求前调用：冷却中抛 RateLimitError；否则等待最小间隔（可读设置）。"""
        interval = resolve_live_order_interval_sec()
        with self._lock:
            self.min_interval_sec = interval
            now = time.monotonic()
            if now < self._cool_until:
                left = self._cool_until - now
                raise RateLimitError(
                    f"{self.name} rate-limit cooldown {left:.1f}s",
                    retry_after=left,
                )
            gap = now - self._last_at
            wait = interval - gap
        if wait > 0:
            time.sleep(wait)
        with self._lock:
            self._last_at = time.monotonic()

    def mark_http(self, status_code: int, retry_after: float | None = None) -> None:
        if status_code not in (418, 429):
            return
        if status_code == 418:
            wait = self.cooldown_418_sec
        else:
            wait = float(retry_after) if retry_after and retry_after > 0 else self.cooldown_429_sec
            wait = max(wait, self.cooldown_429_sec)
        with self._lock:
            self._cool_until = time.monotonic() + wait
        logger.warning("%s HTTP %s → cooldown %.0fs", self.name, status_code, wait)

    def mark_seconds(self, seconds: float) -> None:
        wait = max(1.0, float(seconds))
        with self._lock:
            self._cool_until = max(self._cool_until, time.monotonic() + wait)
        logger.warning("%s cooldown %.0fs (manual)", self.name, wait)


_THROTTLES: dict[str, TradeThrottle] = {}
_THROTTLES_LOCK = threading.Lock()


def get_throttle(name: str, **kwargs: Any) -> TradeThrottle:
    with _THROTTLES_LOCK:
        t = _THROTTLES.get(name)
        if t is None:
            t = TradeThrottle(name, **kwargs)
            _THROTTLES[name] = t
        return t


def is_rate_limit_error(exc: BaseException | str) -> bool:
    if isinstance(exc, RateLimitError):
        return True
    text = str(exc).lower()
    needles = (
        "429",
        "418",
        "rate limit",
        "rate-limit",
        "ratelimit",
        "too many request",
        "cooldown",
        "banned",
        "frequency",
        "请求过于频繁",
        "超出频率",
    )
    return any(n in text for n in needles)


def parse_retry_after_header(headers: Any) -> float | None:
    try:
        raw = headers.get("Retry-After") if headers is not None else None
        if raw is None:
            return None
        return float(raw)
    except (TypeError, ValueError):
        return None


class LiveRetryGate:
    """引擎侧失败退避：避免 half_open / pending / liquidity 每秒砸单。"""

    def __init__(
        self,
        *,
        base_sec: float = 2.0,
        max_sec: float = 60.0,
        rate_limit_min_sec: float = 20.0,
        trip_after: int = 12,
        trip_cooldown_sec: float = 180.0,
    ) -> None:
        self.base_sec = float(base_sec)
        self.max_sec = float(max_sec)
        self.rate_limit_min_sec = float(rate_limit_min_sec)
        self.trip_after = int(trip_after)
        self.trip_cooldown_sec = float(trip_cooldown_sec)
        self._fails: dict[str, int] = {}
        self._next_at: dict[str, float] = {}

    def allow(self, key: str) -> tuple[bool, float]:
        """返回 (可否执行, 剩余等待秒)。"""
        left = max(0.0, self._next_at.get(key, 0.0) - time.monotonic())
        return left <= 0.0, left

    def success(self, key: str) -> None:
        self._fails.pop(key, None)
        self._next_at.pop(key, None)

    def fail(self, key: str, *, rate_limited: bool = False) -> float:
        n = int(self._fails.get(key, 0)) + 1
        self._fails[key] = n
        if rate_limited:
            delay = max(self.rate_limit_min_sec, self.rate_limit_min_sec * (1.5 ** min(n - 1, 4)))
            delay = min(delay, 120.0)
        elif n >= self.trip_after:
            delay = self.trip_cooldown_sec
            logger.error(
                "live retry gate tripped key=%s fails=%s cooldown=%.0fs",
                key,
                n,
                delay,
            )
        else:
            delay = min(self.max_sec, self.base_sec * (2 ** min(n - 1, 5)))
        self._next_at[key] = time.monotonic() + delay
        return delay

    def fails(self, key: str) -> int:
        return int(self._fails.get(key, 0))