From 3460979700091ee40852d4f639bd38fbadb71816 Mon Sep 17 00:00:00 2001
From: DevilXD <4180725+DevilXD@users.noreply.github.com>
Date: Thu, 19 Sep 2024 23:34:54 +0200
Subject: [PATCH] Implement GQL rate limiter

---
 twitch.py |  9 +++++----
 utils.py  | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/twitch.py b/twitch.py
index 677d851..ef39804 100644
--- a/twitch.py
+++ b/twitch.py
@@ -36,6 +36,7 @@ from utils import (
     timestamp,
     create_nonce,
     task_wrapper,
+    RateLimiter,
     AwaitableValue,
     ExponentialBackoff,
 )
@@ -428,6 +429,9 @@ class Twitch:
         self.inventory: list[DropsCampaign] = []
         self._drops: dict[str, TimedDrop] = {}
         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
         self._client_type: ClientInfo = ClientType.ANDROID_APP
         self._session: aiohttp.ClientSession | None = None
@@ -1347,18 +1351,15 @@ class Twitch:
         gql_logger.debug(f"GQL Request: {ops}")
         backoff = ExponentialBackoff(maximum=60)
         for delay in backoff:
-            try:
+            async with self._qgl_limiter:
                 auth_state = await self.get_auth()
                 async with self.request(
                     "POST",
                     "https://gql.twitch.tv/gql",
                     json=ops,
                     headers=auth_state.headers(user_agent=self._client_type.USER_AGENT, gql=True),
-                    invalidate_after=None,  # unused
                 ) as response:
                     response_json: JsonType | list[JsonType] = await response.json()
-            except RequestInvalid:
-                continue
             gql_logger.debug(f"GQL Response: {response_json}")
             orig_response = response_json
             if isinstance(response_json, list):
diff --git a/utils.py b/utils.py
index f5a0248..ca0eea9 100644
--- a/utils.py
+++ b/utils.py
@@ -327,6 +327,51 @@ class ExponentialBackoff:
         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]):
     def __init__(self):
         self._value: _T