Browse Source

Handle gateway rate limits by using a rate limiter.

With the new chunking changes this will become necessary and we don't
want to disconnect from having too many outwards requests.
pull/5849/head
Rapptz 5 years ago
parent
commit
e7db9cfdf7
  1. 31
      discord/gateway.py

31
discord/gateway.py

@ -66,6 +66,31 @@ class WebSocketClosure(Exception):
EventListener = namedtuple('EventListener', 'predicate event result future') EventListener = namedtuple('EventListener', 'predicate event result future')
class GatewayRatelimiter:
def __init__(self, count=120, per=60.0):
self.max = count
self.remaining = count
self.window = 0.0
self.per = per
def get_delay(self):
current = time.time()
if current > self.window + self.per:
self.remaining = self.max
if self.remaining == self.max:
self.window = current
if self.remaining == 0:
return self.per - (current - self.window)
self.remaining -= 1
if self.remaining == 0:
self.window = current
return 0.0
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)
@ -240,6 +265,7 @@ class DiscordWebSocket:
self._zlib = zlib.decompressobj() self._zlib = zlib.decompressobj()
self._buffer = bytearray() self._buffer = bytearray()
self._close_code = None self._close_code = None
self._rate_limiter = GatewayRatelimiter()
@property @property
def open(self): def open(self):
@ -532,6 +558,11 @@ 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()
if delay:
log.warning('WebSocket is ratelimited, waiting %.2f seconds', 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)

Loading…
Cancel
Save