feat: add simple ratelimit control

khl/requester.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Union, List
+from typing import Union, List, Dict
 
 from aiohttp import ClientSession
 
@@ -15,9 +15,11 @@
 class HTTPRequester:
     """wrap raw requests, handle boilerplate param filling works"""
 
-    def __init__(self, cert: Cert):
+    def __init__(self, cert: Cert, ratelimit: bool = True):
         self._cert = cert
         self._cs: Union[ClientSession, None] = None
+        self._ratelimit = self.RateLimit()
+        self.ratelimit = ratelimit
 
     def __del__(self):
         if self._cs is not None:
@@ -28,7 +30,14 @@ async def request(self, method: str, route: str, **params) -> Union[dict, list,
         headers = params.pop('headers', {})
         params['headers'] = headers
 
-        log.debug(f'{method} {route}: req: {params}')  # token is excluded
+        if self.ratelimit:
+            bucket = self._ratelimit.get_bucket(route)
+            delay = self._ratelimit.get_delay(bucket)
+            log.debug(f'{method} {route}: req: {params}, ratelimit: {bucket} {delay}s')  # token is excluded
+            await asyncio.sleep(delay)
+        else:
+            log.debug(f'{method} {route}: req: {params}')  # token is excluded
+
         headers['Authorization'] = f'Bot {self._cert.token}'
         if self._cs is None:  # lazy init
             self._cs = ClientSession()
@@ -40,6 +49,15 @@ async def request(self, method: str, route: str, **params) -> Union[dict, list,
                 rsp = rsp['data']
             else:
                 rsp = await res.read()
+
+            if self.ratelimit and 'X-Rate-Limit-Limit' in res.headers:
+                bucket = res.headers['X-Rate-Limit-Bucket']
+                remaining = int(res.headers['X-Rate-Limit-Remaining'])
+                reset = int(res.headers['X-Rate-Limit-Reset'])
+                self._ratelimit.push_api_bucket_mapping(route, bucket)
+                self._ratelimit.update_ratelimit(bucket, remaining, reset)
+                log.debug(f'{method} {route}: rsp ratelimit: {bucket} {remaining} {reset}s')
+
             log.debug(f'{method} {route}: rsp: {rsp}')
             return rsp
 
@@ -108,3 +126,62 @@ def __init__(self, method, route, params, err_code, err_message):
 
         def __str__(self):
             return f"Requesting '{self.method} {self.route}' failed with {self.err_code}: {self.err_message}"
+
+    class RateLimit:
+        """rate limit control"""
+
+        def __init__(self):
+            self._ratelimit_info: Dict[str, HTTPRequester.RateLimit.RateLimitData] = {}
+            self._api_bucket_mapping: Dict[str, str] = {}
+
+        def push_api_bucket_mapping(self, api: str, bucket: str):
+            """
+            when finished request, associate bucket that api returned with api route
+            to avoid that bucket and api router are not the same
+            """
+
+            api = api.lower()
+            bucket = bucket.lower()
+
+            if api not in self._api_bucket_mapping:
+                self._api_bucket_mapping[api] = bucket
+
+        def get_bucket(self, api: str):
+            """get bucket name by api route"""
+
+            api = api.lower()
+            if api not in self._api_bucket_mapping:
+                return api
+            return self._api_bucket_mapping[api]
+
+        def update_ratelimit(self, bucket: str, remaining: int, reset: int):
+            """update rate limit info"""
+
+            bucket = bucket.lower()
+            if bucket not in self._ratelimit_info:
+                self._ratelimit_info[bucket] = self.RateLimitData(remaining, reset)
+            else:
+                self._ratelimit_info[bucket].remaining = remaining
+                self._ratelimit_info[bucket].reset = reset
+
+        def get_delay(self, bucket: str) -> float:
+            """get request delay time, milliseconds"""
+
+            bucket = bucket.lower()
+            if bucket not in self._ratelimit_info:
+                return 0
+
+            if self._ratelimit_info[bucket].reset == 0:
+                return 0
+
+            if self._ratelimit_info[bucket].remaining == 0:
+                return self._ratelimit_info[bucket].reset * 1000
+
+            delay = self._ratelimit_info[bucket].reset / self._ratelimit_info[bucket].remaining
+            return int(delay * 1000) / 1000  # keep three-digit decimal
+
+        class RateLimitData:
+            """to save single bucket rate limit"""
+            def __init__(self, remaining: int = 120, reset: int = 0):
+                self.remaining = remaining
+                self.reset = reset