fix: add heartbeat for the websocket server (#2404)

Author: knagaitsev

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

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);
     });
   }

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);
   });
 });

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",
 ]
 `;