From 17ead282850add7e0268fecee75e0ca3043de84d Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Fri, 16 Apr 2021 17:07:56 -0700
Subject: [PATCH] fix(server): disconnect ws clients on server close (#6215)

---
 src/remote/playwrightServer.ts | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/remote/playwrightServer.ts b/src/remote/playwrightServer.ts
index 1ee420b456..57d35c6a15 100644
--- a/src/remote/playwrightServer.ts
+++ b/src/remote/playwrightServer.ts
@@ -33,7 +33,7 @@ export interface PlaywrightServerDelegate {
 }
 
 export class PlaywrightServer {
-  private _server: http.Server | undefined;
+  private _wsServer: ws.Server | undefined;
   private _clientsCount = 0;
   private _delegate: PlaywrightServerDelegate;
 
@@ -74,11 +74,10 @@ export class PlaywrightServer {
       });
     });
 
-    this._server = server;
     debugLog('Listening at ' + wsEndpoint);
 
-    const wsServer = new ws.Server({ server: this._server, path });
-    wsServer.on('connection', async socket => {
+    this._wsServer = new ws.Server({ server, path });
+    this._wsServer.on('connection', async socket => {
       if (this._clientsCount && !this._delegate.allowMultipleClients) {
         socket.close();
         return;
@@ -117,10 +116,12 @@ export class PlaywrightServer {
   }
 
   async close() {
-    if (!this._server)
+    if (!this._wsServer)
       return;
     debugLog('Closing server');
-    await new Promise(f => this._server!.close(f));
+    // First disconnect all remaining clients.
+    await new Promise(f => this._wsServer!.close(f));
+    await new Promise(f => this._wsServer!.options.server!.close(f));
     await this._delegate.onClose();
   }
 }