siloqy/prod/bingx/rate_limits.py

from __future__ import annotations

from dataclasses import dataclass
from time import monotonic_ns


@dataclass(frozen=True)
class BingxRateLimitSnapshot:
    rest_remaining: int | None = None
    rest_reset_ms: int | None = None
    ws_listenkey_ops: int = 0
    ws_listenkey_window_ns: int = 0


@dataclass(frozen=True)
class BingxCircuitBreakerSnapshot:
    failure_count: int = 0
    open_until_ns: int = 0
    last_delay_ms: int = 0

    @property
    def is_open(self) -> bool:
        return monotonic_ns() < self.open_until_ns


class BingxRateLimitTracker:
    def __init__(self) -> None:
        self._rest_remaining: int | None = None
        self._rest_reset_ms: int | None = None
        self._ws_listenkey_ops = 0
        self._ws_window_start_ns = monotonic_ns()

    def update_rest_headers(self, headers: dict[str, str]) -> None:
        remain = headers.get("x-ratelimit-requests-remain")
        reset = headers.get("x-ratelimit-requests-expire")
        self._rest_remaining = int(remain) if remain is not None and str(remain).isdigit() else self._rest_remaining
        self._rest_reset_ms = int(reset) if reset is not None and str(reset).isdigit() else self._rest_reset_ms

    def count_ws_listenkey_op(self) -> None:
        now = monotonic_ns()
        if now - self._ws_window_start_ns > 1_000_000_000:
            self._ws_window_start_ns = now
            self._ws_listenkey_ops = 0
        self._ws_listenkey_ops += 1

    def snapshot(self) -> BingxRateLimitSnapshot:
        return BingxRateLimitSnapshot(
            rest_remaining=self._rest_remaining,
            rest_reset_ms=self._rest_reset_ms,
            ws_listenkey_ops=self._ws_listenkey_ops,
            ws_listenkey_window_ns=monotonic_ns() - self._ws_window_start_ns,
        )


class BingxCircuitBreaker:
    def __init__(
        self,
        *,
        failure_threshold: int = 3,
        base_backoff_ms: int = 250,
        max_backoff_ms: int = 2_000,
    ) -> None:
        self._failure_threshold = max(1, int(failure_threshold))
        self._base_backoff_ms = max(1, int(base_backoff_ms))
        self._max_backoff_ms = max(self._base_backoff_ms, int(max_backoff_ms))
        self._failure_count = 0
        self._open_until_ns = 0
        self._last_delay_ms = 0

    async def wait_if_open(self) -> None:
        remaining = self.open_remaining_secs()
        if remaining > 0:
            from asyncio import sleep

            await sleep(remaining)

    def open_remaining_secs(self) -> float:
        remaining_ns = self._open_until_ns - monotonic_ns()
        return max(0.0, remaining_ns / 1_000_000_000)

    def snapshot(self) -> BingxCircuitBreakerSnapshot:
        return BingxCircuitBreakerSnapshot(
            failure_count=self._failure_count,
            open_until_ns=self._open_until_ns,
            last_delay_ms=self._last_delay_ms,
        )

    def record_success(self) -> None:
        self._failure_count = 0
        self._open_until_ns = 0
        self._last_delay_ms = 0

    def record_failure(self, *, rate_limited: bool = False, retry_after_ms: int | None = None) -> float:
        now_ns = monotonic_ns()
        if rate_limited and retry_after_ms is not None and retry_after_ms > 0:
            self._failure_count = self._failure_threshold
            self._last_delay_ms = retry_after_ms
            self._open_until_ns = max(self._open_until_ns, now_ns + retry_after_ms * 1_000_000)
            return retry_after_ms / 1000.0

        self._failure_count += 1
        delay_ms = min(self._base_backoff_ms * (2 ** (self._failure_count - 1)), self._max_backoff_ms)
        self._last_delay_ms = delay_ms
        if self._failure_count >= self._failure_threshold:
            self._open_until_ns = max(self._open_until_ns, now_ns + delay_ms * 1_000_000)
        return delay_ms / 1000.0