refactor!: rewrite http ratelimiting (#29)

* refactor!: rewrite http ratelimiting

* docs: update according to new changes

* refactor: bucket should not be changed if hashes are same

Author: EmmmaTech

discatcore/errors.py

@@ -8,7 +8,6 @@
 __all__ = (
     "DisCatCoreException",
     "HTTPException",
-    "BucketMigrated",
     "UnsupportedAPIVersionWarning",
     "GatewayReconnect",
 )
@@ -81,13 +80,6 @@ def __init__(
         super().__init__(format.format(response.status, response.reason, self.code, self.text))  # type: ignore
 
 
-class BucketMigrated(DisCatCoreException):
-    """Represents an internal exception for when a bucket migrates."""
-
-    def __init__(self, discord_hash: str) -> None:
-        super().__init__(f"This bucket has been migrated to a bucket located at {discord_hash}")
-
-
 class UnsupportedAPIVersionWarning(Warning):
     """Represents a warning for unsupported API versions."""
 

discatcore/http/client.py

@@ -12,7 +12,7 @@
 import discord_typings as dt
 
 from .. import __version__
-from ..errors import BucketMigrated, HTTPException, UnsupportedAPIVersionWarning
+from ..errors import HTTPException, UnsupportedAPIVersionWarning
 from ..file import BasicFile
 from ..types import Unset, UnsetOr
 from ..utils.json import dumps, loads
@@ -226,7 +226,8 @@ async def request(
                 (like Create Guild Sticker).
 
         Returns:
-            If this route returns t.Any content, it will be processed and returned.
+            If this route returns any content, it will be processed and returned. If the request fails after 5 tries,
+            Unset will be returned instead.
         """
         self._request_id += 1
         rid = self._request_id
@@ -236,6 +237,7 @@ async def request(
         query_params = _filter_dict_for_unset(query_params or {})
         max_tries = 5
         headers: dict[str, str] = self.default_headers
+        bucket_hash: t.Optional[str] = None
 
         if reason:
             headers["X-Audit-Log-Reason"] = _urlquote(reason, safe="/ ")
@@ -249,12 +251,13 @@ async def request(
         if data.multipart_content is not Unset:
             kwargs["data"] = data.multipart_content
 
-        bucket = self._ratelimiter.get_bucket(route.bucket)
+        for try_ in range(max_tries):
+            bucket = self._ratelimiter.get_bucket((route.bucket, bucket_hash))
 
-        for tries in range(max_tries):
             async with self._ratelimiter.global_bucket:
+                _log.debug("REQUEST:%d The global ratelimit bucket has been acquired!", rid)
                 async with bucket:
-                    _log.debug("REQUEST:%d All ratelimit buckets have been acquired!", rid)
+                    _log.debug("REQUEST:%d The route ratelimit bucket has been acquired!", rid)
 
                     response = await self._session.request(
                         route.method,
@@ -264,24 +267,28 @@ async def request(
                         **kwargs,
                     )
                     _log.debug(
-                        "REQUEST:%d Made request to %s with method %s.",
+                        "REQUEST:%d Made request to %s with method %s and got status code %d.",
                         rid,
                         f"{self._api_url}{url}",
                         route.method,
+                        response.status,
                     )
 
-                    URL_BUCKET = bucket.bucket is None
-                    bucket.update_info(response)
-                    await bucket.acquire()
-
-                    if URL_BUCKET and bucket.bucket is not None:
-                        try:
-                            self._ratelimiter.migrate_bucket(route.bucket, bucket.bucket)
-                        except BucketMigrated:
-                            bucket = self._ratelimiter.get_bucket(route.bucket)
+                    bucket_hash = response.headers.get("X-RateLimit-Bucket")
+                    if bucket_hash is not None and bucket_hash != bucket.bucket:
+                        _log.debug(
+                            "REQUEST:%d Migrating from bucket (%s, %s) to bucket (%s, %s).",
+                            rid,
+                            route.bucket,
+                            bucket.bucket,
+                            route.bucket,
+                            bucket_hash,
+                        )
+                        bucket = self._ratelimiter.get_bucket((route.bucket, bucket_hash))
 
                     # Everything is ok
                     if 200 <= response.status < 300:
+                        bucket.update_info(response)
                         return await self._text_or_json(response)
 
                     # Ratelimited
@@ -302,13 +309,23 @@ async def request(
                             )
                             self._ratelimiter.global_bucket.lock_for(retry_after)
                             await self._ratelimiter.global_bucket.acquire()
+                        else:
+                            _log.info(
+                                "REQUEST:%d All requests with bucket (%s, %s) have hit a ratelimit! Retrying in %f.",
+                                rid,
+                                route.bucket,
+                                bucket.bucket,
+                                retry_after,
+                            )
+                            bucket.lock_for(retry_after)
+                            await bucket.acquire()
 
                         _log.info("REQUEST:%d Ratelimit is over. Continuing with the request.", rid)
                         continue
 
                     # Specific Server Errors, retry after some time
                     if response.status in {500, 502, 504}:
-                        wait_time = 1 + tries * 2
+                        wait_time = 1 + try_ * 2
                         _log.info("REQUEST:%d Got a server error! Retrying in %d.", rid, wait_time)
                         await asyncio.sleep(wait_time)
                         continue
@@ -317,9 +334,14 @@ async def request(
                     if response.status >= 400:
                         raise HTTPException(response, await self._text_or_json(response))
 
-        raise RuntimeError(
-            f'REQUEST:{rid} Tried sending request to "{url}" with method {route.method} {max_tries} times.'
+        _log.error(
+            'REQUEST:%d Tried sending request to "%s" with method %s %d times.',
+            rid,
+            url,
+            route.method,
+            max_tries,
         )
+        return Unset
 
     async def get_gateway_bot(self) -> dt.GetGatewayBotData:
         """Fetches the gateway information from the Discord API.

discatcore/http/ratelimiter.py

@@ -1,16 +1,12 @@
 # SPDX-License-Identifier: MIT
 
-import logging
 import typing as t
 from datetime import datetime, timezone
 
 from aiohttp import ClientResponse
 
-from ..errors import BucketMigrated
 from ..utils.ratelimit import BurstRatelimiter, ManualRatelimiter
 
-_log = logging.getLogger(__name__)
-
 __all__ = (
     "Bucket",
     "Ratelimiter",
@@ -31,7 +27,6 @@ class Bucket(BurstRatelimiter):
         "reset",
         "bucket",
         "_first_update",
-        "_migrated",
     )
 
     def __init__(self) -> None:
@@ -40,7 +35,6 @@ def __init__(self) -> None:
         self.reset: t.Optional[datetime] = None
         self.bucket: t.Optional[str] = None
         self._first_update: bool = True
-        self._migrated: bool = False
 
     def update_info(self, response: ClientResponse) -> None:
         """Updates the bucket's underlying information via the new headers.
@@ -86,76 +80,30 @@ def update_info(self, response: ClientResponse) -> None:
         if self._first_update:
             self._first_update = False
 
-        if self.reset_after is not None and self.remaining == 0 and not self.is_locked():
-            self.lock_for(self.reset_after)
-
-    def migrate_to(self, discord_hash: str) -> t.NoReturn:
-        """Migrates this bucket to a new one provided by the Discord API.
-
-        Raises:
-            BucketMigrated: An internal exception for the request function to change buckets.
-        """
-        self._migrated = True
-        raise BucketMigrated(discord_hash)
-
-    @property
-    def migrated(self) -> bool:
-        return self._migrated
-
 
 class Ratelimiter:
     """Represents the global ratelimiter.
 
     Attributes:
-        discord_buckets (dict[str, Bucket]): A mapping that maps Discord hashes to bucket objects.
-        url_buckets (dict[str, Bucket]): A mapping that maps psuedo-buckets to bucket objects.
-            This is primarily used by new requests that do not know their Discord hashes.
-        url_to_discord_hash (dict[str, str]): A mapping that maps psuedo-buckets to Discord hashes.
-            This is primarily set by requests that have just discovered their Discord hashes.
-        global_bucket (ManualRatelimiter): The global bucket. Used for requests that involve global 429s.
+        buckets: A mapping of route urls and bucket hashes to buckets.
+        global_bucket: The global bucket. Used for requests that involve global 429s.
     """
 
-    __slots__ = ("discord_buckets", "url_buckets", "url_to_discord_hash", "global_bucket")
+    __slots__ = ("buckets", "global_bucket")
 
     def __init__(self) -> None:
-        self.discord_buckets: dict[str, Bucket] = {}
-        self.url_buckets: dict[str, Bucket] = {}
-        self.url_to_discord_hash: dict[str, str] = {}
+        self.buckets: dict[tuple[str, t.Optional[str]], Bucket] = {}
         self.global_bucket = ManualRatelimiter()
 
-    def get_bucket(self, url: str) -> Bucket:
-        """Gets a bucket object from the providing url.
+    def get_bucket(self, key: tuple[str, t.Optional[str]]) -> Bucket:
+        """Gets a bucket object with the provided key.
 
         Args:
-            url (str): The url to grab the bucket with.
-                This can either be a pseudo-bucket or an actual Discord hash.
+            key: The key to grab the bucket with. This key is in the format (route_url, bucket_hash).
         """
-        if url not in self.url_to_discord_hash:
-            # this serves as a temporary bucket until further ratelimiting info is provided
-            # or since some routes have no ratelimiting, they have to backfire to this bucket instead
+        if key not in self.buckets:
             new_bucket = Bucket()
-            self.url_buckets[url] = new_bucket
-            return new_bucket
-
-        discord_hash = self.url_to_discord_hash[url]
-        return self.discord_buckets[discord_hash]
-
-    def migrate_bucket(self, url: str, discord_hash: str) -> t.NoReturn:
-        """Migrates a bucket object from a pseudo-bucket to a Discord hash.
-        Not only will this call :meth:`Bucket.migrate_to`, but the mappings will be updated.
-
-        Args:
-            url (str): The pseudo-bucket of the bucket to migrate.
-            discord_hash (str): The Discord hash to migrate the bucket to.
-        """
-        self.url_to_discord_hash[url] = discord_hash
-
-        cur_bucket = self.url_buckets.get(url)
-
-        if not cur_bucket:
-            raise BucketMigrated(discord_hash)
-
-        self.discord_buckets[discord_hash] = cur_bucket
-        del self.url_buckets[url]
+            self.buckets[key] = new_bucket
+            new_bucket.bucket = key[1]
 
-        cur_bucket.migrate_to(discord_hash)
+        return self.buckets[key]

discatcore/http/route.py

@@ -36,24 +36,31 @@ def __init__(self, method: str, url: str, **params: t.Any) -> None:
         self.url: str = url
 
         # top-level resource parameters
-        self.guild_id: t.Optional[dt.Snowflake] = params.get("guild_id")
-        self.channel_id: t.Optional[dt.Snowflake] = params.get("channel_id")
-        self.webhook_id: t.Optional[dt.Snowflake] = params.get("webhook_id")
-        self.webhook_token: t.Optional[str] = params.get("webhook_token")
+        self.guild_id: t.Optional[dt.Snowflake] = params.pop("guild_id", None)
+        self.channel_id: t.Optional[dt.Snowflake] = params.pop("channel_id", None)
+        self.webhook_id: t.Optional[dt.Snowflake] = params.pop("webhook_id", None)
+        self.webhook_token: t.Optional[str] = params.pop("webhook_token", None)
 
     @property
     def endpoint(self) -> str:
         """The formatted url for this route."""
-        return self.url.format_map({k: _urlquote(str(v)) for k, v in self.params.items()})
+        top_level_params = {
+            k: getattr(self, k)
+            for k in ("guild_id", "channel_id", "webhook_id", "webhook_token")
+            if getattr(self, k) is not None
+        }
+        other_params = {k: _urlquote(str(v)) for k, v in self.params.items()}
+
+        return self.url.format_map({**top_level_params, **other_params})
 
     @property
     def bucket(self) -> str:
-        """The pseudo-bucket that represents this route. This is generated via the method, raw url and top level parameters."""
+        """The pseudo-bucket that represents this route. This is generated with the method and top level parameters filled into the raw url."""
         top_level_params = {
             k: getattr(self, k)
             for k in ("guild_id", "channel_id", "webhook_id", "webhook_token")
             if getattr(self, k) is not None
         }
-        other_params = {k: None for k in self.params.keys() if k not in top_level_params.keys()}
+        other_params = {k: None for k in self.params.keys()}
 
         return f"{self.method}:{self.url.format_map({**top_level_params, **other_params})}"

discatcore/utils/ratelimit.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: MIT
 
 import asyncio
+import logging
 import typing as t
 
 __all__ = (
@@ -9,6 +10,8 @@
     "BurstRatelimiter",
 )
 
+_log = logging.getLogger(__name__)
+
 
 class BaseRatelimiter:
     """The base class for all ratelimiters. Locking algorithms are up to the subclassed Ratelimiter."""
@@ -35,7 +38,7 @@ async def __aexit__(self, *args: t.Any) -> None:
 
 
 class ManualRatelimiter(BaseRatelimiter):
-    """A simple ratelimiter that simply locks at the command of t.Anything."""
+    """A simple ratelimiter that simply locks at the command of anything."""
 
     async def _unlock(self, delay: float) -> None:
         await asyncio.sleep(delay)
@@ -74,6 +77,7 @@ def __init__(self) -> None:
 
     async def acquire(self) -> None:
         if self.reset_after is not None and self.remaining == 0 and not self.is_locked():
+            _log.info("Auto-locking for %f seconds.", self.reset_after)
             self.lock_for(self.reset_after)
 
         return await super().acquire()