From 2410fdf8d29543daf792ff8fb0e34fa2fc97bdf3 Mon Sep 17 00:00:00 2001
From: Gus Caplan <gustavcaplan@gmail.com>
Date: Sun, 25 Dec 2016 10:28:36 -0600
Subject: [PATCH] fix heartbeats once and for all (#1016)

---
 src/client/Client.js                          |  2 +-
 src/client/ClientManager.js                   | 17 +----------------
 src/client/websocket/WebSocketManager.js      | 19 +++++++++++++++++++
 .../packets/WebSocketPacketManager.js         |  1 +
 .../packets/handlers/GuildMembersChunk.js     |  2 ++
 .../websocket/packets/handlers/Ready.js       |  2 +-
 6 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/src/client/Client.js b/src/client/Client.js
index 4629a38fe..4fdac08db 100644
--- a/src/client/Client.js
+++ b/src/client/Client.js
@@ -416,7 +416,7 @@ class Client extends EventEmitter {
   _pong(startTime) {
     this.pings.unshift(Date.now() - startTime);
     if (this.pings.length > 3) this.pings.length = 3;
-    this.clearTimeout(this._ackTimeout);
+    this.ws.lastHeartbeatAck = true;
   }
 
   _setPresence(id, presence) {
diff --git a/src/client/ClientManager.js b/src/client/ClientManager.js
index ca582668b..0cfbbfdf4 100644
--- a/src/client/ClientManager.js
+++ b/src/client/ClientManager.js
@@ -48,22 +48,7 @@ class ClientManager {
    * @param {number} time The interval in milliseconds at which heartbeat packets should be sent
    */
   setupKeepAlive(time) {
-    this.heartbeatInterval = this.client.setInterval(this.ping.bind(this), time);
-  }
-
-  ping() {
-    this.client.emit('debug', 'Sending heartbeat');
-    this.client._pingTimestamp = Date.now();
-    this.client.ws.send({
-      op: Constants.OPCodes.HEARTBEAT,
-      d: this.client.ws.sequence,
-    }, true);
-
-    const lastPing = this.client.ping;
-
-    this.client._ackTimeout = this.client.setTimeout(() => {
-      this.client.ws.ws.close(1005);
-    }, lastPing ? lastPing * 20 : 20e3);
+    this.heartbeatInterval = this.client.setInterval(() => this.client.ws.heartbeat(true), time);
   }
 
   destroy() {
diff --git a/src/client/websocket/WebSocketManager.js b/src/client/websocket/WebSocketManager.js
index 04772cf41..985105713 100644
--- a/src/client/websocket/WebSocketManager.js
+++ b/src/client/websocket/WebSocketManager.js
@@ -88,6 +88,8 @@ class WebSocketManager extends EventEmitter {
     for (const event of client.options.disabledEvents) this.disabledEvents[event] = true;
 
     this.first = true;
+
+    this.lastHeartbeatAck = true;
   }
 
   /**
@@ -122,6 +124,22 @@ class WebSocketManager extends EventEmitter {
     }
   }
 
+  heartbeat(normal) {
+    if (normal && !this.lastHeartbeatAck) {
+      this.ws.close(1007);
+      return;
+    }
+
+    this.client.emit('debug', 'Sending heartbeat');
+    this.client._pingTimestamp = Date.now();
+    this.client.ws.send({
+      op: Constants.OPCodes.HEARTBEAT,
+      d: this.sequence,
+    }, true);
+
+    this.lastHeartbeatAck = false;
+  }
+
   /**
    * Sends a packet to the gateway
    * @param {Object} data An object that can be JSON stringified
@@ -167,6 +185,7 @@ class WebSocketManager extends EventEmitter {
    */
   eventOpen() {
     this.client.emit('debug', 'Connection to gateway opened');
+    this.lastHeartbeatAck = true;
     if (this.status === Constants.Status.RECONNECTING) this._sendResume();
     else this._sendNewIdentify();
   }
diff --git a/src/client/websocket/packets/WebSocketPacketManager.js b/src/client/websocket/packets/WebSocketPacketManager.js
index 099579912..3bfb8af44 100644
--- a/src/client/websocket/packets/WebSocketPacketManager.js
+++ b/src/client/websocket/packets/WebSocketPacketManager.js
@@ -84,6 +84,7 @@ class WebSocketPacketManager {
 
     if (packet.op === Constants.OPCodes.HEARTBEAT_ACK) {
       this.ws.client._pong(this.ws.client._pingTimestamp);
+      this.ws.lastHeartbeatAck = true;
       this.ws.client.emit('debug', 'Heartbeat acknowledged');
     }
 
diff --git a/src/client/websocket/packets/handlers/GuildMembersChunk.js b/src/client/websocket/packets/handlers/GuildMembersChunk.js
index 5dda5ad27..e9312e291 100644
--- a/src/client/websocket/packets/handlers/GuildMembersChunk.js
+++ b/src/client/websocket/packets/handlers/GuildMembersChunk.js
@@ -16,6 +16,8 @@ class GuildMembersChunkHandler extends AbstractHandler {
 
     guild._checkChunks();
     client.emit(Constants.Events.GUILD_MEMBERS_CHUNK, members);
+
+    client.ws.lastHeartbeatAck = true;
   }
 }
 
diff --git a/src/client/websocket/packets/handlers/Ready.js b/src/client/websocket/packets/handlers/Ready.js
index 4f49980bb..10bc6b257 100644
--- a/src/client/websocket/packets/handlers/Ready.js
+++ b/src/client/websocket/packets/handlers/Ready.js
@@ -7,7 +7,7 @@ class ReadyHandler extends AbstractHandler {
     const client = this.packetManager.client;
     const data = packet.d;
 
-    client.manager.ping();
+    client.ws.heartbeat();
 
     const clientUser = new ClientUser(client, data.user);
     client.user = clientUser;