ForcePilot/backend/package/yuxi/channels/adapters/wechat/rate_limiter.py

from __future__ import annotations

import asyncio
import time


class TokenBucketRateLimiter:
    def __init__(self, rate: int, per: float = 60.0):
        self._rate = rate
        self._per = per
        self._tokens = float(rate)
        self._last_refill = time.monotonic()
        self._lock = asyncio.Lock()

    async def acquire(self) -> bool:
        async with self._lock:
            return self._try_acquire()

    def acquire_sync(self) -> bool:
        return self._try_acquire_sync()

    def _try_acquire(self) -> bool:
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._tokens = min(self._rate, self._tokens + elapsed * (self._rate / self._per))
        self._last_refill = now

        if self._tokens >= 1.0:
            self._tokens -= 1.0
            return True
        return False

    def _try_acquire_sync(self) -> bool:
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._tokens = min(self._rate, self._tokens + elapsed * (self._rate / self._per))
        self._last_refill = now

        if self._tokens >= 1.0:
            self._tokens -= 1.0
            return True
        return False

    @property
    def available_tokens(self) -> float:
        return self._tokens

    @property
    def rate(self) -> int:
        return self._rate

    @property
    def period(self) -> float:
        return self._per