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.

1 files changed, 31 insertions, 0 deletions

diff --git a/discord/gateway.py b/discord/gateway.py
index ab595566..9a605cd6 100644
--- a/discord/gateway.py
+++ b/discord/gateway.py
@@ -66,6 +66,31 @@ class WebSocketClosure(Exception):
 
 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):
     def __init__(self, *args, **kwargs):
         ws = kwargs.pop('ws', None)
@@ -240,6 +265,7 @@ class DiscordWebSocket:
         self._zlib = zlib.decompressobj()
         self._buffer = bytearray()
         self._close_code = None
+        self._rate_limiter = GatewayRatelimiter()
 
     @property
     def open(self):
@@ -532,6 +558,11 @@ class DiscordWebSocket:
                 raise ConnectionClosed(self.socket, shard_id=self.shard_id, code=code) from None
 
     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)
         await self.socket.send_str(data)