Handle HEARTBEAT_ACK

1 files changed, 21 insertions, 1 deletions

diff --git a/discord/gateway.py b/discord/gateway.py
index d21a6356..386a21c6 100644
--- a/discord/gateway.py
+++ b/discord/gateway.py
@@ -25,6 +25,7 @@ DEALINGS IN THE SOFTWARE.
 """
 
 import sys
+import time
 import websockets
 import asyncio
 import aiohttp
@@ -64,9 +65,23 @@ class KeepAliveHandler(threading.Thread):
         self.daemon = True
         self.msg = 'Keeping websocket alive with sequence {0[d]}'
         self._stop_ev = threading.Event()
+        self._last_ack = time.time()
 
     def run(self):
         while not self._stop_ev.wait(self.interval):
+            if self._last_ack + 2 * self.interval < time.time():
+                log.warn("We have stopped responding to the gateway.")
+                coro = self.ws.close(1006)
+                f = compat.run_coroutine_threadsafe(coro, loop=self.ws.loop)
+
+                try:
+                    f.result()
+                except:
+                    pass
+                finally:
+                    self.stop()
+                    return
+
             data = self.get_payload()
             log.debug(self.msg.format(data))
             coro = self.ws.send_as_json(data)
@@ -86,12 +101,16 @@ class KeepAliveHandler(threading.Thread):
     def stop(self):
         self._stop_ev.set()
 
+    def ack(self):
+        self._last_ack = time.time()
+
 class VoiceKeepAliveHandler(KeepAliveHandler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.msg = 'Keeping voice websocket alive with timestamp {0[d]}'
 
     def get_payload(self):
+        self.ack()
         return {
             'op': self.ws.HEARTBEAT,
             'd': int(time.time() * 1000)
@@ -302,7 +321,8 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
             raise ReconnectWebSocket()
 
         if op == self.HEARTBEAT_ACK:
-            return # disable noisy logging for now
+            self._keep_alive.ack()
+            return
 
         if op == self.HEARTBEAT:
             beat = self._keep_alive.get_payload()