From 1a7c8273de5a5b164c63c9919950babd7ecfaadb Mon Sep 17 00:00:00 2001
From: Loonride <knagaitsev1@gmail.com>
Date: Fri, 31 Jan 2020 07:27:37 -0600
Subject: [PATCH] fix: add heartbeat for the websocket server  (#2404)

---
 lib/Server.js                                 |   1 +
 lib/servers/WebsocketServer.js                |  15 +++
 test/server/servers/WebsocketServer.test.js   | 120 +++++++++++-------
 .../WebsocketServer.test.js.snap              |   5 +-
 4 files changed, 91 insertions(+), 50 deletions(-)

diff --git a/lib/Server.js b/lib/Server.js
index dff653aceb..1e8cf2d1b4 100644
--- a/lib/Server.js
+++ b/lib/Server.js
@@ -71,6 +71,7 @@ class Server {
 
     updateCompiler(this.compiler, this.options);
 
+    this.heartbeatInterval = 30000;
     // this.SocketServerImplementation is a class, so it must be instantiated before use
     this.socketServerImplementation = getSocketServerImplementation(
       this.options
diff --git a/lib/servers/WebsocketServer.js b/lib/servers/WebsocketServer.js
index 8424b6d382..34002c675a 100644
--- a/lib/servers/WebsocketServer.js
+++ b/lib/servers/WebsocketServer.js
@@ -27,6 +27,17 @@ module.exports = class WebsocketServer extends BaseServer {
     this.wsServer.on('error', (err) => {
       this.server.log.error(err.message);
     });
+
+    const noop = () => {};
+
+    setInterval(() => {
+      this.wsServer.clients.forEach((ws) => {
+        if (ws.isAlive === false) return ws.terminate();
+
+        ws.isAlive = false;
+        ws.ping(noop);
+      });
+    }, this.server.heartbeatInterval);
   }
 
   send(connection, message) {
@@ -45,6 +56,10 @@ module.exports = class WebsocketServer extends BaseServer {
   // f should be passed the resulting connection and the connection headers
   onConnection(f) {
     this.wsServer.on('connection', (connection, req) => {
+      connection.isAlive = true;
+      connection.on('pong', () => {
+        connection.isAlive = true;
+      });
       f(connection, req.headers);
     });
   }
diff --git a/test/server/servers/WebsocketServer.test.js b/test/server/servers/WebsocketServer.test.js
index e8e59a786b..252a031f67 100644
--- a/test/server/servers/WebsocketServer.test.js
+++ b/test/server/servers/WebsocketServer.test.js
@@ -7,87 +7,111 @@ const WebsocketServer = require('../../../lib/servers/WebsocketServer');
 const port = require('../../ports-map').WebsocketServer;
 
 describe('WebsocketServer', () => {
+  let server;
   let socketServer;
   let listeningApp;
 
-  beforeAll((done) => {
+  beforeEach((done) => {
     // eslint-disable-next-line new-cap
     const app = new express();
 
     listeningApp = http.createServer(app);
     listeningApp.listen(port, 'localhost', () => {
-      const server = {
+      server = {
         log: {
           error: () => {},
           debug: () => {},
         },
         sockPath: '/ws-server',
         listeningApp,
+        heartbeatInterval: 800,
       };
-
       socketServer = new WebsocketServer(server);
-
       done();
     });
   });
 
-  describe('server', () => {
-    it('should recieve connection, send message, and close client', (done) => {
-      const data = [];
-
-      let headers;
-      socketServer.onConnection((connection, h) => {
-        headers = h;
-        data.push('open');
-        socketServer.send(connection, 'hello world');
-        setTimeout(() => {
-          // the server closes the connection with the client
-          socketServer.close(connection);
-        }, 1000);
-      });
+  it('should recieve connection, send message, and close client', (done) => {
+    const data = [];
 
-      // eslint-disable-next-line new-cap
-      const client = new ws(`http://localhost:${port}/ws-server`);
+    let headers;
+    socketServer.onConnection((connection, h) => {
+      headers = h;
+      data.push('open');
+      socketServer.send(connection, 'hello world');
+      setTimeout(() => {
+        // the server closes the connection with the client
+        socketServer.close(connection);
+      }, 1000);
+    });
 
-      client.onmessage = (e) => {
-        data.push(e.data);
-      };
+    // eslint-disable-next-line new-cap
+    const client = new ws(`http://localhost:${port}/ws-server`);
 
-      client.onclose = () => {
-        data.push('close');
-      };
+    client.onmessage = (e) => {
+      data.push(e.data);
+    };
 
-      setTimeout(() => {
-        expect(headers.host).toMatchSnapshot();
-        expect(data).toMatchSnapshot();
-        done();
-      }, 3000);
+    client.onclose = () => {
+      data.push('close');
+    };
+
+    // the heartbeat interval was shortened greatly above
+    // so that the client is quickly pinged
+    client.on('ping', () => {
+      data.push('ping');
     });
 
-    it('should receive client close event', (done) => {
-      let receivedClientClose = false;
-      socketServer.onConnection((connection) => {
-        socketServer.onConnectionClose(connection, () => {
-          receivedClientClose = true;
-        });
+    setTimeout(() => {
+      expect(headers.host).toMatchSnapshot();
+      expect(data).toMatchSnapshot();
+      done();
+    }, 3000);
+  });
+
+  it('should receive client close event', (done) => {
+    let receivedClientClose = false;
+    socketServer.onConnection((connection) => {
+      socketServer.onConnectionClose(connection, () => {
+        receivedClientClose = true;
       });
+    });
 
-      // eslint-disable-next-line new-cap
-      const client = new ws(`http://localhost:${port}/ws-server`);
+    // eslint-disable-next-line new-cap
+    const client = new ws(`http://localhost:${port}/ws-server`);
 
-      setTimeout(() => {
-        // the client closes itself, the server does not close it
-        client.close();
-      }, 1000);
+    setTimeout(() => {
+      // the client closes itself, the server does not close it
+      client.close();
+    }, 1000);
 
-      setTimeout(() => {
-        expect(receivedClientClose).toBeTruthy();
-        done();
-      }, 3000);
+    setTimeout(() => {
+      expect(receivedClientClose).toBeTruthy();
+      done();
+    }, 3000);
+  });
+
+  it('should terminate a client that is not alive', (done) => {
+    let receivedClientClose = false;
+    socketServer.onConnection((connection) => {
+      // this makes the server think the client did not respond
+      // to a ping in time, so the server will terminate it
+      connection.isAlive = false;
+      socketServer.onConnectionClose(connection, () => {
+        receivedClientClose = true;
+      });
     });
+
+    // eslint-disable-next-line new-cap, no-unused-vars
+    const client = new ws(`http://localhost:${port}/ws-server`);
+
+    setTimeout(() => {
+      expect(receivedClientClose).toBeTruthy();
+      done();
+    }, 3000);
   });
 
-  afterAll((done) => {
+  afterEach((done) => {
     listeningApp.close(done);
   });
 });
diff --git a/test/server/servers/__snapshots__/WebsocketServer.test.js.snap b/test/server/servers/__snapshots__/WebsocketServer.test.js.snap
index 2722e7f517..6e3c1a0e4f 100644
--- a/test/server/servers/__snapshots__/WebsocketServer.test.js.snap
+++ b/test/server/servers/__snapshots__/WebsocketServer.test.js.snap
@@ -1,11 +1,12 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`WebsocketServer server should recieve connection, send message, and close client 1`] = `"localhost:8131"`;
+exports[`WebsocketServer should recieve connection, send message, and close client 1`] = `"localhost:8131"`;
 
-exports[`WebsocketServer server should recieve connection, send message, and close client 2`] = `
+exports[`WebsocketServer should recieve connection, send message, and close client 2`] = `
 Array [
   "open",
   "hello world",
+  "ping",
   "close",
 ]
 `;