Refactor handling of WebSocket connections

Send data to all connected clients. Fixes: webpack#1248
lpinca · Jan 7, 2018 · c8cb2a8 · c8cb2a8
1 parent 2463030
commit c8cb2a8
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 80 deletions.
diff --git a/lib/DevServer.js b/lib/DevServer.js
@@ -15,7 +15,7 @@ const OptionsValidationError = require('./OptionsValidationError');
 const optionsSchema = require('./schemas/options.json');
 const plugins = require('./plugins');
 const { http, https } = require('./server');
-const { sendStats } = require('./util');
+const { send, sendStats } = require('./util');
 
 const defaults = {
   clientLogLevel: 'info',
@@ -57,25 +57,15 @@ module.exports = class DevServer extends EventEmitter {
 
     this.log = log;
     this.options = options;
-
-    plugins(compiler, this);
-
     this.app = app.bind(this)(options);
-    this.devMiddleware = devMiddleware(compiler, options);
-    this.watchers = [];
     this.websocketProxies = [];
 
-
-    features(this);
-
     if (options.https) {
       this.server = https(this.app, log, options);
     } else {
       this.server = http(this.app);
     }
 
-    killable(this.server);
-
     // Proxy websockets without the initial http request
     // https://github.com/chimurai/http-proxy-middleware#external-websocket-upgrade
     this.websocketProxies.forEach((proxy) => {
@@ -91,6 +81,14 @@ module.exports = class DevServer extends EventEmitter {
     this.wss.on('listening', () => {
       log.info('WebSocket Server Attached and Listening');
     });
+
+    plugins(compiler, this);
+
+    this.devMiddleware = devMiddleware(compiler, options);
+    this.watchers = [];
+
+    features(this);
+    killable(this.server);
   }
 
   close(callback) {
@@ -128,47 +126,10 @@ module.exports = class DevServer extends EventEmitter {
       wss.on('connection', (ws) => {
         this.emit('ws-connected', ws);
 
-        this.socket = ws;
-        this.socket.payload = payload;
-
-        function payload(type, data) {
-          return JSON.stringify({ type, data });
-        }
-
-        const originalSend = ws.send;
-
-        ws.send = function send(...args) {
-          if (ws.readyState !== WebSocket.OPEN) {
-            return;
-          }
-
-          const cb = function cb(error) {
-            // we'll occasionally get an Error('not open'); here
-            if (error) {
-              // wait a half second and try again
-              setTimeout(() => {
-                log.debug('ws.send: retrying:', args);
-                originalSend.apply(ws, args);
-              }, 500);
-            }
-          };
-
-          args.push(cb);
-          log.debug('ws.send:', args);
-          originalSend.apply(ws, args);
-
-          // emit the same events we're sending through websockets, so interested
-          // parties can tune in via API without connecitng a client socket
-          const [payld] = args;
-          if (payld && payld.type) {
-            this.emit(`ws-${payld.type}`, payld.data);
-          }
-        };
-
-        ws.send(payload('options', this.options));
+        send(ws, this, 'options', this.options);
 
         if (this.stats) {
-          sendStats(this.socket, log, {
+          sendStats(ws, this, {
             force: true,
             stats: this.stats.toJson(this.options.clientStats)
           });

diff --git a/lib/features.js b/lib/features.js
@@ -6,9 +6,10 @@ const express = require('express');
 const historyApiFallback = require('connect-history-api-fallback');
 const serveIndex = require('serve-index');
 const proxy = require('./proxy');
+const { send } = require('./util');
 
 module.exports = function features(devServer) {
-  const { app, devMiddleware, options, socket } = devServer;
+  const { app, devMiddleware, options, wss } = devServer;
 
   function setHeaders(req, res, next) {
     if (options.headers) {
@@ -69,7 +70,7 @@ module.exports = function features(devServer) {
 
     devServer.emit('watch', path);
     const watcher = chokidar.watch(path, opts).on('change', () => {
-      socket.send(socket.payload('content-changed'));
+      wss.clients.forEach(ws => send(ws, devServer, 'content-changed'));
     });
 
     devServer.watchers.push(watcher);

diff --git a/lib/plugins.js b/lib/plugins.js
@@ -1,17 +1,14 @@
 'use strict';
 
 const webpack = require('webpack');
-const { sendStats } = require('./util');
+const { send, sendStats } = require('./util');
 
 module.exports = function plugins(compiler, devServer) {
-  const { log, options } = devServer;
+  const { log, options, wss } = devServer;
 
   function invalid() {
     devServer.emit('compiler-invalid', compiler);
-    const { socket } = devServer;
-    if (socket) {
-      socket.send(socket.payload('invalid'));
-    }
+    wss.clients.forEach(ws => send(ws, devServer, 'invalid'));
   }
 
   const definePlugin = new webpack.DefinePlugin({
@@ -25,8 +22,6 @@ module.exports = function plugins(compiler, devServer) {
 
   if (options.progress) {
     const progressPlugin = new webpack.ProgressPlugin((percent, message, info) => {
-      const { socket } = devServer;
-
       percent = Math.floor(percent * 100);
       devServer.emit('progress', percent);
 
@@ -39,13 +34,12 @@ module.exports = function plugins(compiler, devServer) {
         message = `${message} (${info})`;
       }
 
-      if (socket) {
-        socket.send(socket.payload('progress-update', { percent, message }));
-
+      wss.clients.forEach((ws) => {
+        send(ws, devServer, 'progress-update', { percent, message });
         if (percent === 100) {
-          socket.send(socket.payload('progress-complete', {}));
+          send(ws, devServer, 'progress-complete', {});
         }
-      }
+      });
     });
     compiler.apply(progressPlugin);
   }
@@ -54,9 +48,11 @@ module.exports = function plugins(compiler, devServer) {
   compiler.plugin('invalid', invalid);
   compiler.plugin('done', (stats) => {
     devServer.emit('compiler-done', compiler, stats);
-    sendStats(devServer.socket, log, {
-      force: true,
-      stats: stats.toJson(options.clientStats)
+    wss.clients.forEach((ws) => {
+      sendStats(ws, devServer, {
+        force: true,
+        stats: stats.toJson(options.clientStats)
+      });
     });
     devServer.stats = stats;
   });

diff --git a/lib/util.js b/lib/util.js
@@ -64,35 +64,45 @@ module.exports = {
     });
   },
 
-  sendStats(socket, log, options) {
+  sendStats(socket, devServer, options) {
     const { force, stats } = options;
-    const send = (type, data) => {
-      if (socket) {
-        socket.send(socket.payload(type, data));
-      }
-    };
 
     if (!stats) {
-      log.error('sendStats: options.stats is undefined');
+      devServer.log.error('sendStats: options.stats is undefined');
     }
 
     if (stats.errors && stats.errors.length > 0) {
-      send('errors', stats.errors);
+      this.send(socket, devServer, 'errors', stats.errors);
       return;
     }
 
     if (!force && stats.assets && stats.assets.every(asset => !asset.emitted)) {
-      send('still-ok');
+      this.send(socket, devServer, 'still-ok');
       return;
     }
 
-    send('hash', stats.hash);
+    this.send(socket, devServer, 'hash', stats.hash);
 
     if (stats.warnings.length > 0) {
-      send('warnings', stats.warnings);
+      this.send(socket, devServer, 'warnings', stats.warnings);
     } else {
-      send('ok');
+      this.send(socket, devServer, 'ok');
+    }
+  },
+
+  send(socket, devServer, type, data) {
+    if (socket.readyState !== socket.OPEN) {
+      return;
     }
+
+    const str = JSON.stringify({ type, data });
+
+    devServer.log.debug('ws.send:', str);
+    socket.send(str);
+
+    // emit the same events we're sending through websockets, so interested
+    // parties can tune in via API without connecitng a client socket
+    devServer.emit(`ws-${type}`, data);
   },
 
   validateHost(headers, options) {