From 87c9c95bb87397b201e185b6e93b46dbc557d6e4 Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Fri, 5 Aug 2022 23:19:16 -0400
Subject: [PATCH] Use persistent dictionary for ratelimit information

This should prevent ratelimit information from being cleared too early.
In order to prevent the dictionary from growing to large expired keys
are deleted once they've been deleted.

At present I'm unsure if this change would cause too much CPU pressure.
---
 discord/http.py | 52 +++++++++++++++++++++++++++++++++++--------------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/discord/http.py b/discord/http.py
index e795ccd1e..365970336 100644
--- a/discord/http.py
+++ b/discord/http.py
@@ -47,7 +47,6 @@ from typing import (
 )
 from urllib.parse import quote as _uriquote
 from collections import deque
-import weakref
 import datetime
 
 import aiohttp
@@ -329,6 +328,19 @@ class Ratelimit:
     everything into a single lock queue per route.
     """
 
+    __slots__ = (
+        'limit',
+        'remaining',
+        'outgoing',
+        'reset_after',
+        'expires',
+        'dirty',
+        '_max_ratelimit_timeout',
+        '_loop',
+        '_pending_requests',
+        '_sleeping',
+    )
+
     def __init__(self, max_ratelimit_timeout: Optional[float]) -> None:
         self.limit: int = 1
         self.remaining: int = self.limit
@@ -392,7 +404,6 @@ class Ratelimit:
                     future.set_exception(exception)
                 else:
                     future.set_result(None)
-                self._has_just_awaken = True
                 awaken += 1
 
             if awaken >= count:
@@ -481,12 +492,12 @@ class HTTPClient:
         # Route key -> Bucket hash
         self._bucket_hashes: Dict[str, str] = {}
         # Bucket Hash + Major Parameters -> Rate limit
-        self._buckets: weakref.WeakValueDictionary[str, Ratelimit] = weakref.WeakValueDictionary()
+        # or
         # Route key + Major Parameters -> Rate limit
-        # Used for temporary one shot requests that don't have a bucket hash
-        # While I'd love to use a single mapping for these, doing this would cause the rate limit objects
-        # to inexplicably be evicted from the dictionary before we're done with it
-        self._oneshots: weakref.WeakValueDictionary[str, Ratelimit] = weakref.WeakValueDictionary()
+        # When the key is the latter, it is used for temporary
+        # one shot requests that don't have a bucket hash
+        # When this reaches 256 elements, it will try to evict based off of expiry
+        self._buckets: Dict[str, Ratelimit] = {}
         self._global_over: asyncio.Event = MISSING
         self.token: Optional[str] = None
         self.proxy: Optional[str] = proxy
@@ -517,6 +528,22 @@ class HTTPClient:
 
         return await self.__session.ws_connect(url, **kwargs)
 
+    def _try_clear_expired_ratelimits(self) -> None:
+        if len(self._buckets) < 256:
+            return
+
+        keys = [key for key, bucket in self._buckets.items() if bucket.is_expired()]
+        for key in keys:
+            del self._buckets[key]
+
+    def get_ratelimit(self, key: str) -> Ratelimit:
+        try:
+            value = self._buckets[key]
+        except KeyError:
+            self._buckets[key] = value = Ratelimit(self.max_ratelimit_timeout)
+            self._try_clear_expired_ratelimits()
+        return value
+
     async def request(
         self,
         route: Route,
@@ -533,16 +560,11 @@ class HTTPClient:
         try:
             bucket_hash = self._bucket_hashes[route_key]
         except KeyError:
-            key = route_key + route.major_parameters
-            mapping = self._oneshots
+            key = f'{route_key}:{route.major_parameters}'
         else:
-            key = bucket_hash + route.major_parameters
-            mapping = self._buckets
+            key = f'{bucket_hash}:{route.major_parameters}'
 
-        try:
-            ratelimit = mapping[key]
-        except KeyError:
-            mapping[key] = ratelimit = Ratelimit(self.max_ratelimit_timeout)
+        ratelimit = self.get_ratelimit(key)
 
         # header creation
         headers: Dict[str, str] = {