mirror of
https://github.com/Rapptz/discord.py.git
synced 2025-07-02 00:00:02 +00:00
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.
This commit is contained in:
parent
7d20379bd9
commit
87c9c95bb8
@ -47,7 +47,6 @@ from typing import (
|
|||||||
)
|
)
|
||||||
from urllib.parse import quote as _uriquote
|
from urllib.parse import quote as _uriquote
|
||||||
from collections import deque
|
from collections import deque
|
||||||
import weakref
|
|
||||||
import datetime
|
import datetime
|
||||||
|
|
||||||
import aiohttp
|
import aiohttp
|
||||||
@ -329,6 +328,19 @@ class Ratelimit:
|
|||||||
everything into a single lock queue per route.
|
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:
|
def __init__(self, max_ratelimit_timeout: Optional[float]) -> None:
|
||||||
self.limit: int = 1
|
self.limit: int = 1
|
||||||
self.remaining: int = self.limit
|
self.remaining: int = self.limit
|
||||||
@ -392,7 +404,6 @@ class Ratelimit:
|
|||||||
future.set_exception(exception)
|
future.set_exception(exception)
|
||||||
else:
|
else:
|
||||||
future.set_result(None)
|
future.set_result(None)
|
||||||
self._has_just_awaken = True
|
|
||||||
awaken += 1
|
awaken += 1
|
||||||
|
|
||||||
if awaken >= count:
|
if awaken >= count:
|
||||||
@ -481,12 +492,12 @@ class HTTPClient:
|
|||||||
# Route key -> Bucket hash
|
# Route key -> Bucket hash
|
||||||
self._bucket_hashes: Dict[str, str] = {}
|
self._bucket_hashes: Dict[str, str] = {}
|
||||||
# Bucket Hash + Major Parameters -> Rate limit
|
# Bucket Hash + Major Parameters -> Rate limit
|
||||||
self._buckets: weakref.WeakValueDictionary[str, Ratelimit] = weakref.WeakValueDictionary()
|
# or
|
||||||
# Route key + Major Parameters -> Rate limit
|
# Route key + Major Parameters -> Rate limit
|
||||||
# Used for temporary one shot requests that don't have a bucket hash
|
# When the key is the latter, it is used for temporary
|
||||||
# While I'd love to use a single mapping for these, doing this would cause the rate limit objects
|
# one shot requests that don't have a bucket hash
|
||||||
# to inexplicably be evicted from the dictionary before we're done with it
|
# When this reaches 256 elements, it will try to evict based off of expiry
|
||||||
self._oneshots: weakref.WeakValueDictionary[str, Ratelimit] = weakref.WeakValueDictionary()
|
self._buckets: Dict[str, Ratelimit] = {}
|
||||||
self._global_over: asyncio.Event = MISSING
|
self._global_over: asyncio.Event = MISSING
|
||||||
self.token: Optional[str] = None
|
self.token: Optional[str] = None
|
||||||
self.proxy: Optional[str] = proxy
|
self.proxy: Optional[str] = proxy
|
||||||
@ -517,6 +528,22 @@ class HTTPClient:
|
|||||||
|
|
||||||
return await self.__session.ws_connect(url, **kwargs)
|
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(
|
async def request(
|
||||||
self,
|
self,
|
||||||
route: Route,
|
route: Route,
|
||||||
@ -533,16 +560,11 @@ class HTTPClient:
|
|||||||
try:
|
try:
|
||||||
bucket_hash = self._bucket_hashes[route_key]
|
bucket_hash = self._bucket_hashes[route_key]
|
||||||
except KeyError:
|
except KeyError:
|
||||||
key = route_key + route.major_parameters
|
key = f'{route_key}:{route.major_parameters}'
|
||||||
mapping = self._oneshots
|
|
||||||
else:
|
else:
|
||||||
key = bucket_hash + route.major_parameters
|
key = f'{bucket_hash}:{route.major_parameters}'
|
||||||
mapping = self._buckets
|
|
||||||
|
|
||||||
try:
|
ratelimit = self.get_ratelimit(key)
|
||||||
ratelimit = mapping[key]
|
|
||||||
except KeyError:
|
|
||||||
mapping[key] = ratelimit = Ratelimit(self.max_ratelimit_timeout)
|
|
||||||
|
|
||||||
# header creation
|
# header creation
|
||||||
headers: Dict[str, str] = {
|
headers: Dict[str, str] = {
|
||||||
|
Loading…
x
Reference in New Issue
Block a user