[major] Make WebSocketServer#close() close connections cleanly

Instead of abruptly closing all WebSocket connections, try to close them
cleanly using the 1001 status code.

If the HTTP/S server was created internally, then close it and emit the
`'close'` event when it closes. Otherwise, if client tracking is
enabled, then emit the `'close'` event when the number of connections
goes down to zero. Otherwise, emit it in the next tick.

Refs: #1902

Author: lpinca

doc/ws.md

@@ -209,7 +209,8 @@ added when the `clientTracking` is truthy.
 
 Close the HTTP server if created internally, terminate all clients and call
 callback when done. If an external HTTP server is used via the `server` or
-`noServer` constructor options, it must be closed manually.
+`noServer` constructor options, it must be closed manually. The callback is
+called with an `Error` if the server is already closed.
 
 ### server.handleUpgrade(request, socket, head, callback)
 

lib/websocket-server.js

@@ -111,7 +111,11 @@ class WebSocketServer extends EventEmitter {
     }
 
     if (options.perMessageDeflate === true) options.perMessageDeflate = {};
-    if (options.clientTracking) this.clients = new Set();
+    if (options.clientTracking) {
+      this.clients = new Set();
+      this._shouldEmitClose = false;
+    }
+
     this.options = options;
     this._state = RUNNING;
   }
@@ -151,29 +155,41 @@ class WebSocketServer extends EventEmitter {
     if (this._state === CLOSING) return;
     this._state = CLOSING;
 
-    //
-    // Terminate all associated clients.
-    //
-    if (this.clients) {
-      for (const client of this.clients) client.terminate();
-    }
+    if (this.options.noServer || this.options.server) {
+      if (this._server) {
+        this._removeListeners();
+        this._removeListeners = this._server = null;
+      }
+
+      if (this.clients) {
+        if (!this.clients.size) {
+          process.nextTick(emitClose, this);
+        } else {
+          for (const client of this.clients) client.close(1001);
 
-    const server = this._server;
+          this._shouldEmitClose = true;
+        }
+      } else {
+        process.nextTick(emitClose, this);
+      }
+    } else {
+      const server = this._server;
 
-    if (server) {
       this._removeListeners();
       this._removeListeners = this._server = null;
 
+      if (this.clients) {
+        for (const client of this.clients) client.close(1001);
+      }
+
       //
-      // Close the http server if it was internally created.
+      // The HTTP/S server was created internally. Close it, and rely on its
+      // `'close'` event.
       //
-      if (this.options.port != null) {
-        server.close(emitClose.bind(undefined, this));
-        return;
-      }
+      server.close(() => {
+        emitClose(this);
+      });
     }
-
-    process.nextTick(emitClose, this);
   }
 
   /**
@@ -373,7 +389,13 @@ class WebSocketServer extends EventEmitter {
 
     if (this.clients) {
       this.clients.add(ws);
-      ws.on('close', () => this.clients.delete(ws));
+      ws.on('close', () => {
+        this.clients.delete(ws);
+
+        if (this._shouldEmitClose && !this.clients.size) {
+          process.nextTick(emitClose, this);
+        }
+      });
     }
 
     cb(ws, req);

test/sender.test.js

@@ -266,6 +266,31 @@ describe('Sender', () => {
   });
 
   describe('#close', () => {
+    it('throws an error if the first argument is invalid', () => {
+      const mockSocket = new MockSocket();
+      const sender = new Sender(mockSocket);
+
+      assert.throws(
+        () => sender.close('error'),
+        /^TypeError: First argument must be a valid error code number$/
+      );
+
+      assert.throws(
+        () => sender.close(1004),
+        /^TypeError: First argument must be a valid error code number$/
+      );
+    });
+
+    it('throws an error if the message is greater than 123 bytes', () => {
+      const mockSocket = new MockSocket();
+      const sender = new Sender(mockSocket);
+
+      assert.throws(
+        () => sender.close(1000, 'a'.repeat(124)),
+        /^RangeError: The message must not be greater than 123 bytes$/
+      );
+    });
+
     it('should consume all data before closing', (done) => {
       const perMessageDeflate = new PerMessageDeflate({ threshold: 0 });
 

test/websocket-server.test.js

@@ -13,7 +13,7 @@ const os = require('os');
 
 const Sender = require('../lib/sender');
 const WebSocket = require('..');
-const { NOOP } = require('../lib/constants');
+const { EMPTY_BUFFER, NOOP } = require('../lib/constants');
 
 describe('WebSocketServer', () => {
   describe('#ctor', () => {
@@ -223,13 +223,19 @@ describe('WebSocketServer', () => {
       let closes = 0;
       const wss = new WebSocket.Server({ port: 0 }, () => {
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
-        ws.on('close', () => {
+        ws.on('close', (code, reason) => {
+          assert.strictEqual(code, 1001);
+          assert.deepStrictEqual(reason, EMPTY_BUFFER);
+
           if (++closes === 2) done();
         });
       });
 
       wss.on('connection', (ws) => {
-        ws.on('close', () => {
+        ws.on('close', (code, reason) => {
+          assert.strictEqual(code, 1001);
+          assert.deepStrictEqual(reason, EMPTY_BUFFER);
+
           if (++closes === 2) done();
         });
         wss.close();
@@ -309,6 +315,16 @@ describe('WebSocketServer', () => {
       });
     });
 
+    it("emits the 'close' event if client tracking is disabled", (done) => {
+      const wss = new WebSocket.Server({
+        noServer: true,
+        clientTracking: false
+      });
+
+      wss.on('close', done);
+      wss.close();
+    });
+
     it("emits the 'close' event if the server is already closed", (done) => {
       const wss = new WebSocket.Server({ port: 0 }, () => {
         wss.close(() => {
@@ -683,6 +699,7 @@ describe('WebSocketServer', () => {
 
             socket.once('data', (chunk) => {
               assert.strictEqual(chunk[0], 0x88);
+              socket.destroy();
               wss.close(done);
             });
           });

test/websocket.test.js

@@ -1743,51 +1743,6 @@ describe('WebSocket', () => {
       });
     });
 
-    it('throws an error if the first argument is invalid (1/2)', (done) => {
-      const wss = new WebSocket.Server({ port: 0 }, () => {
-        const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
-
-        ws.on('open', () => {
-          assert.throws(
-            () => ws.close('error'),
-            /^TypeError: First argument must be a valid error code number$/
-          );
-
-          wss.close(done);
-        });
-      });
-    });
-
-    it('throws an error if the first argument is invalid (2/2)', (done) => {
-      const wss = new WebSocket.Server({ port: 0 }, () => {
-        const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
-
-        ws.on('open', () => {
-          assert.throws(
-            () => ws.close(1004),
-            /^TypeError: First argument must be a valid error code number$/
-          );
-
-          wss.close(done);
-        });
-      });
-    });
-
-    it('throws an error if the message is greater than 123 bytes', (done) => {
-      const wss = new WebSocket.Server({ port: 0 }, () => {
-        const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
-
-        ws.on('open', () => {
-          assert.throws(
-            () => ws.close(1000, 'a'.repeat(124)),
-            /^RangeError: The message must not be greater than 123 bytes$/
-          );
-
-          wss.close(done);
-        });
-      });
-    });
-
     it('sends the close status code only when necessary', (done) => {
       let sent;
       const wss = new WebSocket.Server({ port: 0 }, () => {
@@ -2934,10 +2889,10 @@ describe('WebSocket', () => {
       });
 
       it('calls the callback if the socket is closed prematurely', (done) => {
+        const called = [];
         const wss = new WebSocket.Server(
           { perMessageDeflate: true, port: 0 },
           () => {
-            const called = [];
             const ws = new WebSocket(`ws://localhost:${wss.address().port}`, {
               perMessageDeflate: { threshold: 0 }
             });
@@ -2966,15 +2921,15 @@ describe('WebSocket', () => {
                 );
               });
             });
-
-            ws.on('close', () => {
-              assert.deepStrictEqual(called, [1, 2]);
-              wss.close(done);
-            });
           }
         );
 
         wss.on('connection', (ws) => {
+          ws.on('close', () => {
+            assert.deepStrictEqual(called, [1, 2]);
+            wss.close(done);
+          });
+
           ws._socket.end();
         });
       });