chore: merge Connection.{close,didDisconnect} (#9524)

This simplifes cleanup logic.
Also markAsRemote() in gridClient.

packages/playwright-core/src/client/browser.ts

@@ -20,7 +20,7 @@ import { Page } from './page';
 import { ChannelOwner } from './channelOwner';
 import { Events } from './events';
 import { BrowserContextOptions } from './types';
-import { isSafeCloseError } from '../utils/errors';
+import { isSafeCloseError, kBrowserClosedError } from '../utils/errors';
 import * as api from '../../types/types';
 import { CDPSession } from './cdpSession';
 import type { BrowserType } from './browserType';
@@ -110,7 +110,7 @@ export class Browser extends ChannelOwner<channels.BrowserChannel, channels.Brow
     try {
       await this._wrapApiCall(async (channel: channels.BrowserChannel) => {
         if (this._shouldCloseConnectionOnClose)
-          this._connection.close();
+          this._connection.close(kBrowserClosedError);
         else
           await channel.close();
         await this._closedPromise;

packages/playwright-core/src/client/browserType.ts

@@ -133,8 +133,9 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
       let browser: Browser;
       const { pipe } = await channel.connect({ wsEndpoint, headers: params.headers, slowMo: params.slowMo, timeout: params.timeout });
       const closePipe = () => pipe.close().catch(() => {});
-      const connection = new Connection(closePipe);
+      const connection = new Connection();
       connection.markAsRemote();
+      connection.on('close', closePipe);
 
       const onPipeClosed = () => {
         // Emulate all pages, contexts and the browser closing upon disconnect.
@@ -144,15 +145,14 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
           context._onClose();
         }
         browser?._didClose();
-        connection.didDisconnect(kBrowserClosedError);
+        connection.close(kBrowserClosedError);
       };
       pipe.on('closed', onPipeClosed);
       connection.onmessage = message => pipe.send({ message }).catch(onPipeClosed);
 
       pipe.on('message', ({ message }) => {
         try {
-          if (!connection!.isDisconnected())
-            connection!.dispatch(message);
+          connection!.dispatch(message);
         } catch (e) {
           console.error(`Playwright: Connection dispatch error`);
           console.error(e);
@@ -176,10 +176,7 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
           browser._logger = logger;
           browser._shouldCloseConnectionOnClose = true;
           browser._setBrowserType((playwright as any)[browser._name]);
-          browser.on(Events.Browser.Disconnected, () => {
-            playwright._cleanup();
-            closePipe();
-          });
+          browser.on(Events.Browser.Disconnected, closePipe);
           fulfill(browser);
         } catch (e) {
           reject(e);

packages/playwright-core/src/client/connection.ts

@@ -60,14 +60,12 @@ export class Connection extends EventEmitter {
   private _lastId = 0;
   private _callbacks = new Map<number, { resolve: (a: any) => void, reject: (a: Error) => void, stackTrace: ParsedStackTrace }>();
   private _rootObject: Root;
-  private _disconnectedErrorMessage: string | undefined;
-  private _onClose?: () => void;
+  private _closedErrorMessage: string | undefined;
   private _isRemote = false;
 
-  constructor(onClose?: () => void) {
+  constructor() {
     super();
     this._rootObject = new Root(this);
-    this._onClose = onClose;
   }
 
   markAsRemote() {
@@ -91,6 +89,9 @@ export class Connection extends EventEmitter {
   }
 
   async sendMessageToServer(object: ChannelOwner, method: string, params: any, maybeStackTrace: ParsedStackTrace | null): Promise<any> {
+    if (this._closedErrorMessage)
+      throw new Error(this._closedErrorMessage);
+
     const guid = object._guid;
     const stackTrace: ParsedStackTrace = maybeStackTrace || { frameTexts: [], frames: [], apiName: '', allFrames: [] };
     const { frames, apiName } = stackTrace;
@@ -102,8 +103,6 @@ export class Connection extends EventEmitter {
     const metadata: channels.Metadata = { stack: frames, apiName };
     this.onmessage({ ...converted, metadata });
 
-    if (this._disconnectedErrorMessage)
-      throw new Error(this._disconnectedErrorMessage);
     return await new Promise((resolve, reject) => this._callbacks.set(id, { resolve, reject, stackTrace }));
   }
 
@@ -112,6 +111,9 @@ export class Connection extends EventEmitter {
   }
 
   dispatch(message: object) {
+    if (this._closedErrorMessage)
+      return;
+
     const { id, guid, method, params, result, error } = message as any;
     if (id) {
       debugLogger.log('channel:response', message);
@@ -144,21 +146,12 @@ export class Connection extends EventEmitter {
     object._channel.emit(method, object._type === 'JsonPipe' ? params : this._replaceGuidsWithChannels(params));
   }
 
-  close() {
-    if (this._onClose)
-      this._onClose();
-  }
-
-  didDisconnect(errorMessage: string) {
-    this._disconnectedErrorMessage = errorMessage;
+  close(errorMessage: string = 'Connection closed') {
+    this._closedErrorMessage = errorMessage;
     for (const callback of this._callbacks.values())
       callback.reject(new Error(errorMessage));
     this._callbacks.clear();
-    this.emit('disconnect');
-  }
-
-  isDisconnected() {
-    return !!this._disconnectedErrorMessage;
+    this.emit('close');
   }
 
   private _replaceGuidsWithChannels(payload: any): any {

packages/playwright-core/src/client/playwright.ts

@@ -49,7 +49,6 @@ export class Playwright extends ChannelOwner<channels.PlaywrightChannel, channel
   readonly selectors: Selectors;
   readonly request: Fetch;
   readonly errors: { TimeoutError: typeof TimeoutError };
-  private _selectorsOwner: SelectorsOwner;
   private _sockets = new Map<string, net.Socket>();
   private _redirectPortForTest: number | undefined;
 
@@ -67,8 +66,13 @@ export class Playwright extends ChannelOwner<channels.PlaywrightChannel, channel
     this.selectors = sharedSelectors;
     this.errors = { TimeoutError };
 
-    this._selectorsOwner = SelectorsOwner.from(initializer.selectors);
-    this.selectors._addChannel(this._selectorsOwner);
+    const selectorsOwner = SelectorsOwner.from(initializer.selectors);
+    this.selectors._addChannel(selectorsOwner);
+    this._connection.on('close', () => {
+      this.selectors._removeChannel(selectorsOwner);
+      for (const uid of this._sockets.keys())
+        this._onSocksClosed(uid);
+    });
   }
 
   _enablePortForwarding(redirectPortForTest?: number) {
@@ -116,8 +120,4 @@ export class Playwright extends ChannelOwner<channels.PlaywrightChannel, channel
     this._sockets.get(uid)?.destroy();
     this._sockets.delete(uid);
   }
-
-  _cleanup() {
-    this.selectors._removeChannel(this._selectorsOwner);
-  }
 }

packages/playwright-core/src/grid/gridClient.ts

@@ -34,9 +34,10 @@ export class GridClient {
     if (errorText)
       throw errorText;
     const connection = new Connection();
+    connection.markAsRemote();
     connection.onmessage = (message: Object) => ws.send(JSON.stringify(message));
     ws.on('message', message => connection.dispatch(JSON.parse(message.toString())));
-    ws.on('close', (code, reason) => connection.didDisconnect(reason));
+    ws.on('close', (code, reason) => connection.close(reason));
     const playwright = await connection.initializePlaywright();
     playwright._enablePortForwarding();
     return new GridClient(ws, playwright);

packages/playwright-core/src/remote/playwrightClient.ts

@@ -49,9 +49,7 @@ export class PlaywrightClient {
         playwright = await connection.initializePlaywright();
         resolve(new PlaywrightClient(playwright, ws));
       });
-      ws.on('close', () => {
-        playwright?._cleanup();
-      });
+      ws.on('close', (code, reason) => connection.close(reason));
     });
     let timer: NodeJS.Timeout;
     try {

tests/browsertype-connect.spec.ts

@@ -475,7 +475,7 @@ test('should properly disconnect when connection closes from the client side', a
   await disconnectedPromise;
   expect(browser.isConnected()).toBe(false);
 
-  expect((await navigationPromise).message).toContain('has been closed');
+  expect((await navigationPromise).message).toContain('Connection closed');
   expect((await waitForNavigationPromise).message).toContain('Navigation failed because page was closed');
   expect((await page.goto(server.EMPTY_PAGE).catch(e => e)).message).toContain('has been closed');
   expect((await page.waitForNavigation().catch(e => e)).message).toContain('Navigation failed because page was closed');