refactor!: move web socket client to web socket server class

Author: alexander-akait

lib/Server.js

@@ -29,7 +29,6 @@ class Server {
     this.compiler = compiler;
     this.options = options;
     this.logger = this.compiler.getInfrastructureLogger('webpack-dev-server');
-    this.webSocketConnections = [];
     this.staticWatchers = [];
     // Keep track of websocket proxies for external websocket upgrade.
     this.webSocketProxies = [];
@@ -193,11 +192,13 @@ class Server {
         msg = `${msg} (${addInfo})`;
       }
 
-      this.sendMessage(this.webSocketConnections, 'progress-update', {
-        percent,
-        msg,
-        pluginName,
-      });
+      if (this.webSocketServer) {
+        this.sendMessage(this.webSocketServer.clients, 'progress-update', {
+          percent,
+          msg,
+          pluginName,
+        });
+      }
 
       if (this.server) {
         this.server.emit('progress-update', { percent, msg, pluginName });
@@ -212,15 +213,17 @@ class Server {
   }
 
   setupHooks() {
-    // Listening for events
-    const invalidPlugin = () => {
-      this.sendMessage(this.webSocketConnections, 'invalid');
-    };
-
     const addHooks = (compiler) => {
-      compiler.hooks.invalid.tap('webpack-dev-server', invalidPlugin);
+      compiler.hooks.invalid.tap('webpack-dev-server', () => {
+        if (this.webSocketServer) {
+          this.sendMessage(this.webSocketServer.clients, 'invalid');
+        }
+      });
       compiler.hooks.done.tap('webpack-dev-server', (stats) => {
-        this.sendStats(this.webSocketConnections, this.getStats(stats));
+        if (this.webSocketServer) {
+          this.sendStats(this.webSocketServer.clients, this.getStats(stats));
+        }
+
         this.stats = stats;
       });
     };
@@ -673,15 +676,19 @@ class Server {
     // eslint-disable-next-line new-cap
     this.webSocketServer = new this.webSocketServerImplementation(this);
 
-    this.webSocketServer.onConnection((connection, headers) => {
-      if (!connection) {
-        return;
-      }
+    this.webSocketServer.implementation.on('connection', (client, request) => {
+      const headers =
+        // eslint-disable-next-line no-nested-ternary
+        typeof request !== 'undefined'
+          ? request.headers
+          : typeof client.headers !== 'undefined'
+          ? client.headers
+          : // eslint-disable-next-line no-undefined
+            undefined;
 
       if (!headers) {
         this.logger.warn(
-          'webSocketServer implementation must pass headers to the callback of onConnection(f) ' +
-            'via f(connection, headers) in order for clients to pass a headers security check'
+          'webSocketServer implementation must pass headers for the "connection" event'
         );
       }
 
@@ -690,48 +697,34 @@ class Server {
         !this.checkHostHeader(headers) ||
         !this.checkOriginHeader(headers)
       ) {
-        this.sendMessage([connection], 'error', 'Invalid Host/Origin header');
+        this.sendMessage([client], 'error', 'Invalid Host/Origin header');
 
-        this.webSocketServer.closeConnection(connection);
+        client.terminate();
 
         return;
       }
 
-      this.webSocketConnections.push(connection);
-
-      this.webSocketServer.onConnectionClose(connection, () => {
-        const idx = this.webSocketConnections.indexOf(connection);
-
-        if (idx >= 0) {
-          this.webSocketConnections.splice(idx, 1);
-        }
-      });
-
       if (this.options.hot === true || this.options.hot === 'only') {
-        this.sendMessage([connection], 'hot');
+        this.sendMessage([client], 'hot');
       }
 
       if (this.options.liveReload) {
-        this.sendMessage([connection], 'liveReload');
+        this.sendMessage([client], 'liveReload');
       }
 
       if (this.options.client && this.options.client.progress) {
-        this.sendMessage(
-          [connection],
-          'progress',
-          this.options.client.progress
-        );
+        this.sendMessage([client], 'progress', this.options.client.progress);
       }
 
       if (this.options.client && this.options.client.overlay) {
-        this.sendMessage([connection], 'overlay', this.options.client.overlay);
+        this.sendMessage([client], 'overlay', this.options.client.overlay);
       }
 
       if (!this.stats) {
         return;
       }
 
-      this.sendStats([connection], this.getStats(this.stats), true);
+      this.sendStats([client], this.getStats(this.stats), true);
     });
   }
 
@@ -1072,8 +1065,7 @@ class Server {
 
   close(callback) {
     if (this.webSocketServer) {
-      this.webSocketServer.close();
-      this.webSocketConnections = [];
+      this.webSocketServer.implementation.close();
     }
 
     const prom = Promise.all(
@@ -1235,12 +1227,14 @@ class Server {
     return false;
   }
 
-  sendMessage(webSocketConnections, type, data) {
-    webSocketConnections.forEach((webSocketConnection) => {
-      this.webSocketServer.send(
-        webSocketConnection,
-        JSON.stringify({ type, data })
-      );
+  // eslint-disable-next-line class-methods-use-this
+  sendMessage(clients, type, data) {
+    clients.forEach((client) => {
+      // `sockjs` uses `1` to indicate client is ready to accept data
+      // `ws` uses `WebSocket.OPEN`, but it is mean `1` too
+      if (client.readyState === 1) {
+        client.send(JSON.stringify({ type, data }));
+      }
     });
   }
 
@@ -1270,7 +1264,7 @@ class Server {
   }
 
   // Send stats to a socket or multiple sockets
-  sendStats(webSocketConnections, stats, force) {
+  sendStats(clients, stats, force) {
     const shouldEmit =
       !force &&
       stats &&
@@ -1280,23 +1274,23 @@ class Server {
       stats.assets.every((asset) => !asset.emitted);
 
     if (shouldEmit) {
-      this.sendMessage(webSocketConnections, 'still-ok');
+      this.sendMessage(clients, 'still-ok');
 
       return;
     }
 
-    this.sendMessage(webSocketConnections, 'hash', stats.hash);
+    this.sendMessage(clients, 'hash', stats.hash);
 
     if (stats.errors.length > 0 || stats.warnings.length > 0) {
       if (stats.warnings.length > 0) {
-        this.sendMessage(webSocketConnections, 'warnings', stats.warnings);
+        this.sendMessage(clients, 'warnings', stats.warnings);
       }
 
       if (stats.errors.length > 0) {
-        this.sendMessage(webSocketConnections, 'errors', stats.errors);
+        this.sendMessage(clients, 'errors', stats.errors);
       }
     } else {
-      this.sendMessage(webSocketConnections, 'ok');
+      this.sendMessage(clients, 'ok');
     }
   }
 
@@ -1337,7 +1331,13 @@ class Server {
     // disabling refreshing on changing the content
     if (this.options.liveReload) {
       watcher.on('change', (item) => {
-        this.sendMessage(this.webSocketConnections, 'static-changed', item);
+        if (this.webSocketServer) {
+          this.sendMessage(
+            this.webSocketServer.clients,
+            'static-changed',
+            item
+          );
+        }
       });
     }
 

lib/servers/BaseServer.js

@@ -5,5 +5,6 @@
 module.exports = class BaseServer {
   constructor(server) {
     this.server = server;
+    this.clients = new Set();
   }
 };

lib/servers/SockJSServer.js

@@ -61,37 +61,24 @@ module.exports = class SockJSServer extends BaseServer {
       ...this.server.options.webSocketServer.options,
       prefix: getPrefix(this.server.options.webSocketServer.options),
     });
-  }
 
-  send(connection, message) {
-    // prevent cases where the server is trying to send data while connection is closing
-    if (connection.readyState === 1) {
-      connection.write(message);
-    }
-  }
+    this.implementation.on('connection', (client) => {
+      // Implement the the same API as for `ws`
+      client.send = client.write;
 
-  close(callback) {
-    [...this.server.webSocketConnections].forEach((socket) => {
-      this.closeConnection(socket);
-    });
-
-    if (callback) {
-      callback();
-    }
-  }
+      this.clients.add(client);
 
-  closeConnection(connection) {
-    connection.close();
-  }
-
-  // f should be passed the resulting connection and the connection headers
-  onConnection(f) {
-    this.implementation.on('connection', (connection) => {
-      f(connection, connection ? connection.headers : null);
+      client.on('close', () => {
+        this.clients.delete(client);
+      });
     });
-  }
 
-  onConnectionClose(connection, f) {
-    connection.on('close', f);
+    this.implementation.close = () => {
+      for (const client of this.clients) {
+        client.close();
+      }
+
+      this.clients.clear();
+    };
   }
 };

lib/servers/WebsocketServer.js

@@ -7,13 +7,14 @@ const WebSocket = require('ws');
 const BaseServer = require('./BaseServer');
 
 module.exports = class WebsocketServer extends BaseServer {
-  static heartbeatInterval = 30000;
+  static heartbeatInterval = 1000;
 
   constructor(server) {
     super(server);
 
     this.implementation = new WebSocket.Server({
       ...this.server.options.webSocketServer.options,
+      clientTracking: false,
       noServer: true,
     });
 
@@ -31,51 +32,37 @@ module.exports = class WebsocketServer extends BaseServer {
       this.server.logger.error(err.message);
     });
 
-    const noop = () => {};
-
     const interval = setInterval(() => {
-      this.implementation.clients.forEach((socket) => {
-        if (socket.isAlive === false) {
-          return socket.terminate();
+      this.clients.forEach((client) => {
+        if (client.isAlive === false) {
+          client.terminate();
+
+          return;
         }
 
-        socket.isAlive = false;
-        socket.ping(noop);
+        client.isAlive = false;
+        client.ping(() => {});
       });
     }, WebsocketServer.heartbeatInterval);
 
-    this.implementation.on('close', () => {
-      clearInterval(interval);
-    });
-  }
-
-  send(connection, message) {
-    // Prevent cases where the server is trying to send data while connection is closing
-    if (connection.readyState === WebSocket.OPEN) {
-      connection.send(message);
-    }
-  }
+    this.implementation.on('connection', (client) => {
+      this.clients.add(client);
 
-  close(callback) {
-    this.implementation.close(callback);
-  }
+      client.isAlive = true;
 
-  closeConnection(connection) {
-    connection.close();
-  }
+      client.on('pong', () => {
+        client.isAlive = true;
+      });
 
-  // f should be passed the resulting connection and the connection headers
-  onConnection(f) {
-    this.implementation.on('connection', (connection, req) => {
-      connection.isAlive = true;
-      connection.on('pong', () => {
-        connection.isAlive = true;
+      client.on('close', () => {
+        this.clients.delete(client);
       });
-      f(connection, req.headers);
     });
-  }
 
-  onConnectionClose(connection, f) {
-    connection.on('close', f);
+    this.implementation.on('close', () => {
+      clearInterval(interval);
+
+      this.clients.clear();
+    });
   }
 };

test/e2e/web-socket-communication.test.js

@@ -66,7 +66,7 @@ describe('web socket communication', () => {
 
       await new Promise((resolve) => {
         const interval = setInterval(() => {
-          if (server.webSocketConnections.length === 0) {
+          if (server.webSocketServer.clients.size === 0) {
             clearInterval(interval);
             resolve();
           }
@@ -140,7 +140,7 @@ describe('web socket communication', () => {
         }, 200);
       });
 
-      expect(server.webSocketConnections).toHaveLength(0);
+      expect(server.webSocketServer.clients.size).toBe(0);
       expect(consoleMessages.map((message) => message.text())).toMatchSnapshot(
         'console messages'
       );