限制请求开放平台的频率

api/open_live.py

@@ -16,6 +16,7 @@ import tornado.web
 import api.base
 import config
 import services.open_live
+import utils.rate_limit
 import utils.request
 
 logger = logging.getLogger(__name__)
@@ -32,6 +33,8 @@ END_GAME_COMMON_SERVER_URL = COMMON_SERVER_BASE_URL + '/api/internal/open_live/e
 GAME_HEARTBEAT_COMMON_SERVER_URL = COMMON_SERVER_BASE_URL + '/api/internal/open_live/game_heartbeat'
 
 _error_auth_code_cache = cachetools.LRUCache(256)
+# 用于限制请求开放平台的频率
+_open_live_rate_limiter = utils.rate_limit.TokenBucket(5, 9)
 
 
 class TransportError(Exception):
@@ -66,7 +69,7 @@ async def request_open_live_or_common_server(open_live_url, common_server_url, b
         raise
 
 
-async def request_open_live(url, body: dict) -> dict:
+async def request_open_live(url, body: dict, *, ignore_rate_limit=False) -> dict:
     cfg = config.get_config()
     assert cfg.is_open_live_configured
 
@@ -77,6 +80,10 @@ async def request_open_live(url, body: dict) -> dict:
     else:
         auth_code = ''
 
+    # 频率限制，防止触发B站风控被下架
+    if not _open_live_rate_limiter.try_decrease_token() and not ignore_rate_limit:
+        raise BusinessError({'code': 4009, 'message': '接口访问限制', 'request_id': '0', 'data': None})
+
     body_bytes = json.dumps(body).encode('utf-8')
     headers = {
         'x-bili-accesskeyid': cfg.open_live_access_key_id,

services/open_live.py

@@ -73,7 +73,8 @@ async def _flush_game_heartbeat_tasks():
     try:
         res = await api.open_live.request_open_live(
             api.open_live.GAME_BATCH_HEARTBEAT_OPEN_LIVE_URL,
-            {'game_ids': game_ids}
+            {'game_ids': game_ids},
+            ignore_rate_limit=True
         )
         failed_game_ids = res['data']['failed_game_ids']
         if failed_game_ids is None:  # 哪个SB后端给数组传null的

utils/rate_limit.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+import datetime
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class TokenBucket:
+    def __init__(self, tokens_per_sec, max_token_num):
+        self._tokens_per_sec = float(tokens_per_sec)
+        self._max_token_num = float(max_token_num)
+        self._stored_token_num = self._max_token_num
+        self._last_update_time = datetime.datetime.now()
+
+        if self._tokens_per_sec <= 0.0 and self._max_token_num >= 1.0:
+            logger.warning('TokenBucket token_per_sec=%f <= 0, rate has no limit', tokens_per_sec)
+
+    def try_decrease_token(self):
+        if self._tokens_per_sec <= 0.0:
+            # self._max_token_num < 1.0 时完全禁止
+            return self._max_token_num >= 1.0
+
+        cur_time = datetime.datetime.now()
+        last_update_time = min(self._last_update_time, cur_time)  # 防止时钟回拨
+        add_token_num = (cur_time - last_update_time).total_seconds() * self._tokens_per_sec
+        self._stored_token_num = min(self._stored_token_num + add_token_num, self._max_token_num)
+        self._last_update_time = cur_time
+
+        if self._stored_token_num < 1.0:
+            return False
+        self._stored_token_num -= 1.0
+        return True