matthew/enhanced-discord.py
1
0
Fork 0
Code Issues 4 Pull Requests 21 Projects Releases Wiki Activity

Use a lock for the gateway rate limiter.

Browse Source 
This will allow for higher concurrency in AutoSharded situations where
I can mostly "fire and forget" the chunk requests.
This commit is contained in:
Rapptz 2020-09-05 12:20:58 -04:00
parent 5fe998ac19
commit af8fc32329
1 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
discord/gateway.py
View File

@ -73,6 +73,8 @@ class GatewayRatelimiter:
self.remaining = count self.remaining = count
self.window = 0.0 self.window = 0.0
self.per = per self.per = per
self.lock = asyncio.Lock()
self.shard_id = None
def get_delay(self): def get_delay(self):
current = time.time() current = time.time()
@ -92,6 +94,14 @@ class GatewayRatelimiter:
return 0.0 return 0.0
async def block(self):
async with self.lock:
delta = self.get_delay()
if delta:
log.warning('WebSocket in shard ID %s is ratelimited, waiting %.2f seconds', self.shard_id, delta)
await asyncio.sleep(delta)
class KeepAliveHandler(threading.Thread): class KeepAliveHandler(threading.Thread):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
ws = kwargs.pop('ws', None) ws = kwargs.pop('ws', None)
@ -291,6 +301,7 @@ class DiscordWebSocket:
ws.call_hooks = client._connection.call_hooks ws.call_hooks = client._connection.call_hooks
ws._initial_identify = initial ws._initial_identify = initial
ws.shard_id = shard_id ws.shard_id = shard_id
ws._rate_limiter.shard_id = shard_id
ws.shard_count = client._connection.shard_count ws.shard_count = client._connection.shard_count
ws.session_id = session ws.session_id = session
ws.sequence = sequence ws.sequence = sequence
@ -559,11 +570,7 @@ class DiscordWebSocket:
raise ConnectionClosed(self.socket, shard_id=self.shard_id, code=code) from None raise ConnectionClosed(self.socket, shard_id=self.shard_id, code=code) from None
async def send(self, data): async def send(self, data):
delay = self._rate_limiter.get_delay() await self._rate_limiter.block()
if delay:
log.warning('WebSocket in shard ID %s is ratelimited, waiting %.2f seconds', self.shard_id, delay)
await asyncio.sleep(delay)
self._dispatch('socket_raw_send', data) self._dispatch('socket_raw_send', data)
await self.socket.send_str(data) await self.socket.send_str(data)