chore: extract ws server util (#29247)

microsoft · Jan 30, 2024 · aff6cf3 · aff6cf3
1 parent aeafd44
commit aff6cf3
Show file tree

Hide file tree

Showing 29 changed files with 295 additions and 203 deletions.
diff --git a/packages/playwright-core/src/client/channelOwner.ts b/packages/playwright-core/src/client/channelOwner.ts
@@ -17,7 +17,7 @@
 import { EventEmitter } from 'events';
 import type * as channels from '@protocol/channels';
 import { maybeFindValidator, ValidationError, type ValidatorContext } from '../protocol/validator';
-import { debugLogger } from '../common/debugLogger';
+import { debugLogger } from '../utils/debugLogger';
 import type { ExpectZone } from '../utils/stackTrace';
 import { captureRawStack, captureLibraryStackTrace, stringifyStackFrames } from '../utils/stackTrace';
 import { isUnderTest } from '../utils';

diff --git a/packages/playwright-core/src/client/connection.ts b/packages/playwright-core/src/client/connection.ts
@@ -32,7 +32,7 @@ import { Electron, ElectronApplication } from './electron';
 import type * as channels from '@protocol/channels';
 import { Stream } from './stream';
 import { WritableStream } from './writableStream';
-import { debugLogger } from '../common/debugLogger';
+import { debugLogger } from '../utils/debugLogger';
 import { SelectorsOwner } from './selectors';
 import { Android, AndroidSocket, AndroidDevice } from './android';
 import { Artifact } from './artifact';

diff --git a/packages/playwright-core/src/client/harRouter.ts b/packages/playwright-core/src/client/harRouter.ts
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { debugLogger } from '../common/debugLogger';
+import { debugLogger } from '../utils/debugLogger';
 import type { BrowserContext } from './browserContext';
 import type { LocalUtils } from './localUtils';
 import type { Route } from './network';

diff --git a/packages/playwright-core/src/common/socksProxy.ts b/packages/playwright-core/src/common/socksProxy.ts
@@ -17,7 +17,7 @@
 import EventEmitter from 'events';
 import type { AddressInfo } from 'net';
 import net from 'net';
-import { debugLogger } from './debugLogger';
+import { debugLogger } from '../utils/debugLogger';
 import { createSocket } from '../utils/happy-eyeballs';
 import { assert, createGuid,  } from '../utils';
 

diff --git a/packages/playwright-core/src/remote/playwrightConnection.ts b/packages/playwright-core/src/remote/playwrightConnection.ts
@@ -27,7 +27,7 @@ import { AndroidDevice } from '../server/android/android';
 import { DebugControllerDispatcher } from '../server/dispatchers/debugControllerDispatcher';
 import { startProfiling, stopProfiling } from '../utils';
 import { monotonicTime } from '../utils';
-import { debugLogger } from '../common/debugLogger';
+import { debugLogger } from '../utils/debugLogger';
 
 export type ClientType = 'controller' | 'launch-browser' | 'reuse-browser' | 'pre-launched-browser-or-android';
 

diff --git a/packages/playwright-core/src/remote/playwrightServer.ts b/packages/playwright-core/src/remote/playwrightServer.ts
@@ -14,24 +14,18 @@
  * limitations under the License.
  */
 
-import { wsServer } from '../utilsBundle';
-import type { WebSocketServer } from '../utilsBundle';
-import type http from 'http';
 import type { Browser } from '../server/browser';
 import type { Playwright } from '../server/playwright';
 import { createPlaywright } from '../server/playwright';
 import { PlaywrightConnection } from './playwrightConnection';
 import type { ClientType } from './playwrightConnection';
 import type  { LaunchOptions } from '../server/types';
-import { ManualPromise } from '../utils/manualPromise';
+import { Semaphore } from '../utils/semaphore';
 import type { AndroidDevice } from '../server/android/android';
 import type { SocksProxy } from '../common/socksProxy';
-import { debugLogger } from '../common/debugLogger';
-import { createHttpServer, userAgentVersionMatchesErrorMessage } from '../utils';
-import { perMessageDeflate } from '../server/transport';
-
-let lastConnectionId = 0;
-const kConnectionSymbol = Symbol('kConnection');
+import { debugLogger } from '../utils/debugLogger';
+import { userAgentVersionMatchesErrorMessage } from '../utils';
+import { WSServer } from '../utils/wsServer';
 
 type ServerOptions = {
   path: string;
@@ -44,193 +38,94 @@ type ServerOptions = {
 
 export class PlaywrightServer {
   private _preLaunchedPlaywright: Playwright | undefined;
-  private _wsServer: WebSocketServer | undefined;
-  private _server: http.Server | undefined;
   private _options: ServerOptions;
+  private _wsServer: WSServer;
 
   constructor(options: ServerOptions) {
     this._options = options;
     if (options.preLaunchedBrowser)
       this._preLaunchedPlaywright = options.preLaunchedBrowser.attribution.playwright;
     if (options.preLaunchedAndroidDevice)
       this._preLaunchedPlaywright = options.preLaunchedAndroidDevice._android.attribution.playwright;
-  }
-
-  async listen(port: number = 0, hostname?: string): Promise<string> {
-    debugLogger.log('server', `Server started at ${new Date()}`);
-
-    const server = createHttpServer((request: http.IncomingMessage, response: http.ServerResponse) => {
-      if (request.method === 'GET' && request.url === '/json') {
-        response.setHeader('Content-Type', 'application/json');
-        response.end(JSON.stringify({
-          wsEndpointPath: this._options.path,
-        }));
-        return;
-      }
-      response.end('Running');
-    });
-    server.on('error', error => debugLogger.log('server', String(error)));
-    this._server = server;
-
-    const wsEndpoint = await new Promise<string>((resolve, reject) => {
-      server.listen(port, hostname, () => {
-        const address = server.address();
-        if (!address) {
-          reject(new Error('Could not bind server socket'));
-          return;
-        }
-        const wsEndpoint = typeof address === 'string' ? `${address}${this._options.path}` : `ws://${hostname || 'localhost'}:${address.port}${this._options.path}`;
-        resolve(wsEndpoint);
-      }).on('error', reject);
-    });
 
-    debugLogger.log('server', 'Listening at ' + wsEndpoint);
-    this._wsServer = new wsServer({
-      noServer: true,
-      perMessageDeflate,
-    });
     const browserSemaphore = new Semaphore(this._options.maxConnections);
     const controllerSemaphore = new Semaphore(1);
     const reuseBrowserSemaphore = new Semaphore(1);
-    if (process.env.PWTEST_SERVER_WS_HEADERS) {
-      this._wsServer.on('headers', (headers, request) => {
-        headers.push(process.env.PWTEST_SERVER_WS_HEADERS!);
-      });
-    }
-    server.on('upgrade', (request, socket, head) => {
-      const pathname = new URL('http://localhost' + request.url!).pathname;
-      if (pathname !== this._options.path) {
-        socket.write(`HTTP/${request.httpVersion} 400 Bad Request\r\n\r\n`);
-        socket.destroy();
-        return;
-      }
 
-      const uaError = userAgentVersionMatchesErrorMessage(request.headers['user-agent'] || '');
-      if (uaError) {
-        socket.write(`HTTP/${request.httpVersion} 428 Precondition Required\r\n\r\n${uaError}`);
-        socket.destroy();
-        return;
-      }
-
-      this._wsServer?.handleUpgrade(request, socket, head, ws => this._wsServer?.emit('connection', ws, request));
-    });
-    this._wsServer.on('connection', (ws, request) => {
-      debugLogger.log('server', 'Connected client ws.extension=' + ws.extensions);
-      const url = new URL('http://localhost' + (request.url || ''));
-      const browserHeader = request.headers['x-playwright-browser'];
-      const browserName = url.searchParams.get('browser') || (Array.isArray(browserHeader) ? browserHeader[0] : browserHeader) || null;
-      const proxyHeader = request.headers['x-playwright-proxy'];
-      const proxyValue = url.searchParams.get('proxy') || (Array.isArray(proxyHeader) ? proxyHeader[0] : proxyHeader);
-
-      const launchOptionsHeader = request.headers['x-playwright-launch-options'] || '';
-      const launchOptionsHeaderValue = Array.isArray(launchOptionsHeader) ? launchOptionsHeader[0] : launchOptionsHeader;
-      const launchOptionsParam = url.searchParams.get('launch-options');
-      let launchOptions: LaunchOptions = {};
-      try {
-        launchOptions = JSON.parse(launchOptionsParam || launchOptionsHeaderValue);
-      } catch (e) {
-      }
+    this._wsServer = new WSServer({
+      onUpgrade: (request, socket) => {
+        const uaError = userAgentVersionMatchesErrorMessage(request.headers['user-agent'] || '');
+        if (uaError)
+          return { error: `HTTP/${request.httpVersion} 428 Precondition Required\r\n\r\n${uaError}` };
+      },
+
+      onHeaders: headers => {
+        if (process.env.PWTEST_SERVER_WS_HEADERS)
+          headers.push(process.env.PWTEST_SERVER_WS_HEADERS!);
+      },
+
+      onConnection: (request, url, ws, id) => {
+        const browserHeader = request.headers['x-playwright-browser'];
+        const browserName = url.searchParams.get('browser') || (Array.isArray(browserHeader) ? browserHeader[0] : browserHeader) || null;
+        const proxyHeader = request.headers['x-playwright-proxy'];
+        const proxyValue = url.searchParams.get('proxy') || (Array.isArray(proxyHeader) ? proxyHeader[0] : proxyHeader);
+
+        const launchOptionsHeader = request.headers['x-playwright-launch-options'] || '';
+        const launchOptionsHeaderValue = Array.isArray(launchOptionsHeader) ? launchOptionsHeader[0] : launchOptionsHeader;
+        const launchOptionsParam = url.searchParams.get('launch-options');
+        let launchOptions: LaunchOptions = {};
+        try {
+          launchOptions = JSON.parse(launchOptionsParam || launchOptionsHeaderValue);
+        } catch (e) {
+        }
 
-      const id = String(++lastConnectionId);
-      debugLogger.log('server', `[${id}] serving connection: ${request.url}`);
+        // Instantiate playwright for the extension modes.
+        const isExtension = this._options.mode === 'extension';
+        if (isExtension) {
+          if (!this._preLaunchedPlaywright)
+            this._preLaunchedPlaywright = createPlaywright({ sdkLanguage: 'javascript', isServer: true });
+        }
 
-      // Instantiate playwright for the extension modes.
-      const isExtension = this._options.mode === 'extension';
-      if (isExtension) {
-        if (!this._preLaunchedPlaywright)
-          this._preLaunchedPlaywright = createPlaywright({ sdkLanguage: 'javascript', isServer: true });
-      }
+        let clientType: ClientType = 'launch-browser';
+        let semaphore: Semaphore = browserSemaphore;
+        if (isExtension && url.searchParams.has('debug-controller')) {
+          clientType = 'controller';
+          semaphore = controllerSemaphore;
+        } else if (isExtension) {
+          clientType = 'reuse-browser';
+          semaphore = reuseBrowserSemaphore;
+        } else if (this._options.mode === 'launchServer') {
+          clientType = 'pre-launched-browser-or-android';
+          semaphore = browserSemaphore;
+        }
 
-      let clientType: ClientType = 'launch-browser';
-      let semaphore: Semaphore = browserSemaphore;
-      if (isExtension && url.searchParams.has('debug-controller')) {
-        clientType = 'controller';
-        semaphore = controllerSemaphore;
-      } else if (isExtension) {
-        clientType = 'reuse-browser';
-        semaphore = reuseBrowserSemaphore;
-      } else if (this._options.mode === 'launchServer') {
-        clientType = 'pre-launched-browser-or-android';
-        semaphore = browserSemaphore;
+        return new PlaywrightConnection(
+            semaphore.acquire(),
+            clientType, ws,
+            { socksProxyPattern: proxyValue, browserName, launchOptions },
+            {
+              playwright: this._preLaunchedPlaywright,
+              browser: this._options.preLaunchedBrowser,
+              androidDevice: this._options.preLaunchedAndroidDevice,
+              socksProxy: this._options.preLaunchedSocksProxy,
+            },
+            id, () => semaphore.release());
+      },
+
+      onClose: async () => {
+        debugLogger.log('server', 'closing browsers');
+        if (this._preLaunchedPlaywright)
+          await Promise.all(this._preLaunchedPlaywright.allBrowsers().map(browser => browser.close({ reason: 'Playwright Server stopped' })));
+        debugLogger.log('server', 'closed browsers');
       }
-
-      const connection = new PlaywrightConnection(
-          semaphore.acquire(),
-          clientType, ws,
-          { socksProxyPattern: proxyValue, browserName, launchOptions },
-          {
-            playwright: this._preLaunchedPlaywright,
-            browser: this._options.preLaunchedBrowser,
-            androidDevice: this._options.preLaunchedAndroidDevice,
-            socksProxy: this._options.preLaunchedSocksProxy,
-          },
-          id, () => semaphore.release());
-      (ws as any)[kConnectionSymbol] = connection;
     });
-
-    return wsEndpoint;
-  }
-
-  async close() {
-    const server = this._wsServer;
-    if (!server)
-      return;
-    debugLogger.log('server', 'closing websocket server');
-    const waitForClose = new Promise(f => server.close(f));
-    // First disconnect all remaining clients.
-    await Promise.all(Array.from(server.clients).map(async ws => {
-      const connection = (ws as any)[kConnectionSymbol] as PlaywrightConnection | undefined;
-      if (connection)
-        await connection.close();
-      try {
-        ws.terminate();
-      } catch (e) {
-      }
-    }));
-    await waitForClose;
-    debugLogger.log('server', 'closing http server');
-    if (this._server)
-      await new Promise(f => this._server!.close(f));
-    this._wsServer = undefined;
-    this._server = undefined;
-    debugLogger.log('server', 'closed server');
-
-    debugLogger.log('server', 'closing browsers');
-    if (this._preLaunchedPlaywright)
-      await Promise.all(this._preLaunchedPlaywright.allBrowsers().map(browser => browser.close({ reason: 'Playwright Server stopped' })));
-    debugLogger.log('server', 'closed browsers');
   }
-}
 
-export class Semaphore {
-  private _max: number;
-  private _acquired = 0;
-  private _queue: ManualPromise[] = [];
-
-  constructor(max: number) {
-    this._max = max;
-  }
-
-  setMax(max: number) {
-    this._max = max;
-  }
-
-  acquire(): Promise<void> {
-    const lock = new ManualPromise();
-    this._queue.push(lock);
-    this._flush();
-    return lock;
-  }
-
-  release() {
-    --this._acquired;
-    this._flush();
+  async listen(port: number = 0, hostname?: string): Promise<string> {
+    return this._wsServer.listen(port, hostname, this._options.path);
   }
 
-  private _flush() {
-    while (this._acquired < this._max && this._queue.length) {
-      ++this._acquired;
-      this._queue.shift()!.resolve();
-    }
+  async close() {
+    await this._wsServer.close();
   }
 }
diff --git a/packages/playwright-core/src/server/android/android.ts b/packages/playwright-core/src/server/android/android.ts
@@ -30,7 +30,7 @@ import { ProgressController } from '../progress';
 import { CRBrowser } from '../chromium/crBrowser';
 import { helper } from '../helper';
 import { PipeTransport } from '../../protocol/transport';
-import { RecentLogsCollector } from '../../common/debugLogger';
+import { RecentLogsCollector } from '../../utils/debugLogger';
 import { gracefullyCloseSet } from '../../utils/processLauncher';
 import { TimeoutSettings } from '../../common/timeoutSettings';
 import type * as channels from '@protocol/channels';

diff --git a/packages/playwright-core/src/server/browser.ts b/packages/playwright-core/src/server/browser.ts
@@ -21,7 +21,7 @@ import { Page } from './page';
 import { Download } from './download';
 import type { ProxySettings } from './types';
 import type { ChildProcess } from 'child_process';
-import type { RecentLogsCollector } from '../common/debugLogger';
+import type { RecentLogsCollector } from '../utils/debugLogger';
 import type { CallMetadata } from './instrumentation';
 import { SdkObject } from './instrumentation';
 import { Artifact } from './artifact';

diff --git a/packages/playwright-core/src/server/browserType.ts b/packages/playwright-core/src/server/browserType.ts
@@ -35,7 +35,7 @@ import { DEFAULT_TIMEOUT, TimeoutSettings } from '../common/timeoutSettings';
 import { debugMode } from '../utils';
 import { existsAsync } from '../utils/fileUtils';
 import { helper } from './helper';
-import { RecentLogsCollector } from '../common/debugLogger';
+import { RecentLogsCollector } from '../utils/debugLogger';
 import type { CallMetadata } from './instrumentation';
 import { SdkObject } from './instrumentation';
 import { ManualPromise } from '../utils/manualPromise';

diff --git a/packages/playwright-core/src/server/chromium/chromium.ts b/packages/playwright-core/src/server/chromium/chromium.ts
@@ -37,7 +37,7 @@ import { getUserAgent } from '../../utils/userAgent';
 import { wrapInASCIIBox } from '../../utils/ascii';
 import { debugMode, headersArrayToObject, headersObjectToArray, } from '../../utils';
 import { removeFolders } from '../../utils/fileUtils';
-import { RecentLogsCollector } from '../../common/debugLogger';
+import { RecentLogsCollector } from '../../utils/debugLogger';
 import type { Progress } from '../progress';
 import { ProgressController } from '../progress';
 import { TimeoutSettings } from '../../common/timeoutSettings';

diff --git a/packages/playwright-core/src/server/chromium/crConnection.ts b/packages/playwright-core/src/server/chromium/crConnection.ts
@@ -19,8 +19,8 @@ import { type RegisteredListener, assert, eventsHelper } from '../../utils';
 import type { ConnectionTransport, ProtocolRequest, ProtocolResponse } from '../transport';
 import type { Protocol } from './protocol';
 import { EventEmitter } from 'events';
-import type { RecentLogsCollector } from '../../common/debugLogger';
-import { debugLogger } from '../../common/debugLogger';
+import type { RecentLogsCollector } from '../../utils/debugLogger';
+import { debugLogger } from '../../utils/debugLogger';
 import type { ProtocolLogger } from '../types';
 import { helper } from '../helper';
 import { ProtocolError } from '../protocolError';

diff --git a/packages/playwright-core/src/server/electron/electron.ts b/packages/playwright-core/src/server/electron/electron.ts
@@ -40,7 +40,7 @@ import type { BrowserOptions, BrowserProcess } from '../browser';
 import type { Playwright } from '../playwright';
 import type * as childProcess from 'child_process';
 import * as readline from 'readline';
-import { RecentLogsCollector } from '../../common/debugLogger';
+import { RecentLogsCollector } from '../../utils/debugLogger';
 import { serverSideCallMetadata, SdkObject } from '../instrumentation';
 import type * as channels from '@protocol/channels';