Implement GQL rate limiter

This commit is contained in:
DevilXD committed 2024-09-19 23:34:54 +02:00
1 parent d68271068d
commit 3460979700
2 files changed
+50 -4

No files matched your search

+5 -4
View File
@@ -36,6 +36,7 @@ from utils import (
timestamp, timestamp,
create_nonce, create_nonce,
task_wrapper, task_wrapper,
RateLimiter,
AwaitableValue, AwaitableValue,
ExponentialBackoff, ExponentialBackoff,
) )
@@ -428,6 +429,9 @@ class Twitch:
self.inventory: list[DropsCampaign] = [] self.inventory: list[DropsCampaign] = []
self._drops: dict[str, TimedDrop] = {} self._drops: dict[str, TimedDrop] = {}
self._mnt_triggers: deque[datetime] = deque() self._mnt_triggers: deque[datetime] = deque()
# NOTE: GQL is pretty volatile and breaks everything if one runs into their rate limit.
# Do not modify the default, safe values.
self._qgl_limiter = RateLimiter(capacity=20, window=1)
# Client type, session and auth # Client type, session and auth
self._client_type: ClientInfo = ClientType.ANDROID_APP self._client_type: ClientInfo = ClientType.ANDROID_APP
self._session: aiohttp.ClientSession | None = None self._session: aiohttp.ClientSession | None = None
@@ -1347,18 +1351,15 @@ class Twitch:
gql_logger.debug(f"GQL Request: {ops}") gql_logger.debug(f"GQL Request: {ops}")
backoff = ExponentialBackoff(maximum=60) backoff = ExponentialBackoff(maximum=60)
for delay in backoff: for delay in backoff:
try: async with self._qgl_limiter:
auth_state = await self.get_auth() auth_state = await self.get_auth()
async with self.request( async with self.request(
"POST", "POST",
"https://gql.twitch.tv/gql", "https://gql.twitch.tv/gql",
json=ops, json=ops,
headers=auth_state.headers(user_agent=self._client_type.USER_AGENT, gql=True), headers=auth_state.headers(user_agent=self._client_type.USER_AGENT, gql=True),
invalidate_after=None, # unused
) as response: ) as response:
response_json: JsonType | list[JsonType] = await response.json() response_json: JsonType | list[JsonType] = await response.json()
except RequestInvalid:
continue
gql_logger.debug(f"GQL Response: {response_json}") gql_logger.debug(f"GQL Response: {response_json}")
orig_response = response_json orig_response = response_json
if isinstance(response_json, list): if isinstance(response_json, list):
+45
View File
@@ -327,6 +327,51 @@ class ExponentialBackoff:
self.steps = 0 self.steps = 0
class RateLimiter:
def __init__(self, *, capacity: int, window: int):
self.total: int = 0
self.concurrent: int = 0
self.window: int = window
self.capacity: int = capacity
self._reset_task: asyncio.Task[None] | None = None
self._cond: asyncio.Condition = asyncio.Condition()
def __repr__(self) -> str:
return f"{self.__class__.__name__}(capacity={self.total}/{self.capacity})"
def __del__(self) -> None:
if self._reset_task is not None:
self._reset_task.cancel()
def _can_proceed(self) -> bool:
return max(self.total, self.concurrent) < self.capacity
async def __aenter__(self):
async with self._cond:
await self._cond.wait_for(self._can_proceed)
self.total += 1
self.concurrent += 1
if self._reset_task is None:
self._reset_task = asyncio.create_task(self._rtask())
async def __aexit__(self, exc_type, exc, tb):
self.concurrent -= 1
async with self._cond:
self._cond.notify(self.capacity - self.concurrent)
async def _reset(self) -> None:
if self._reset_task is not None:
self._reset_task = None
async with self._cond:
self.total = 0
if self.concurrent < self.capacity:
self._cond.notify(self.capacity - self.concurrent)
async def _rtask(self) -> None:
await asyncio.sleep(self.window)
await self._reset()
class AwaitableValue(Generic[_T]): class AwaitableValue(Generic[_T]):
def __init__(self): def __init__(self):
self._value: _T self._value: _T