chore: align client side instrumentations (#9771)

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

@@ -37,6 +37,7 @@ import { Tracing } from './tracing';
 import type { BrowserType } from './browserType';
 import { Artifact } from './artifact';
 import { FetchRequest } from './fetch';
+import { createInstrumentation } from './clientInstrumentation';
 
 export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel, channels.BrowserContextInitializer> implements api.BrowserContext {
   _pages = new Set<Page>();
@@ -64,7 +65,7 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel,
   }
 
   constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.BrowserContextInitializer) {
-    super(parent, type, guid, initializer);
+    super(parent, type, guid, initializer, createInstrumentation());
     if (parent instanceof Browser)
       this._browser = parent;
     this._isChromium = this._browser?._name === 'chromium';

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

@@ -20,8 +20,9 @@ import { createScheme, ValidationError, Validator } from '../protocol/validator'
 import { debugLogger } from '../utils/debugLogger';
 import { captureStackTrace, ParsedStackTrace } from '../utils/stackTrace';
 import { isUnderTest } from '../utils/utils';
+import { ClientInstrumentation } from './clientInstrumentation';
 import type { Connection } from './connection';
-import type { ClientSideInstrumentation, Logger } from './types';
+import type { Logger } from './types';
 
 export abstract class ChannelOwner<T extends channels.Channel = channels.Channel, Initializer = {}> extends EventEmitter {
   readonly _connection: Connection;
@@ -33,15 +34,16 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
   readonly _channel: T;
   readonly _initializer: Initializer;
   _logger: Logger | undefined;
-  _csi: ClientSideInstrumentation | undefined;
+  _instrumentation: ClientInstrumentation | undefined;
 
-  constructor(parent: ChannelOwner | Connection, type: string, guid: string, initializer: Initializer) {
+  constructor(parent: ChannelOwner | Connection, type: string, guid: string, initializer: Initializer, instrumentation?: ClientInstrumentation) {
     super();
     this.setMaxListeners(0);
     this._connection = parent instanceof ChannelOwner ? parent._connection : parent;
     this._type = type;
     this._guid = guid;
     this._parent = parent instanceof ChannelOwner ? parent : undefined;
+    this._instrumentation = instrumentation || this._parent?._instrumentation;
 
     this._connection._objects.set(guid, this);
     if (this._parent) {
@@ -72,7 +74,7 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
     };
   }
 
-  private _createChannel(base: Object, stackTrace: ParsedStackTrace | null, csi?: ClientSideInstrumentation, callCookie?: { userObject: any }): T {
+  private _createChannel(base: Object, stackTrace: ParsedStackTrace | null, csi?: ClientInstrumentation, callCookie?: any): T {
     const channel = new Proxy(base, {
       get: (obj: any, prop) => {
         if (prop === 'debugScopeState')
@@ -82,7 +84,7 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
           if (validator) {
             return (params: any) => {
               if (callCookie && csi) {
-                callCookie.userObject = csi.onApiCallBegin(renderCallWithParams(stackTrace!.apiName!, params), stackTrace).userObject;
+                csi.onApiCallBegin(renderCallWithParams(stackTrace!.apiName!, params), stackTrace, callCookie);
                 csi = undefined;
               }
               return this._connection.sendMessageToServer(this, prop, validator(params, ''), stackTrace);
@@ -101,16 +103,12 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
     const stackTrace = captureStackTrace();
     const { apiName, frameTexts } = stackTrace;
 
-    let ancestorWithCSI: ChannelOwner<any> = this;
-    while (!ancestorWithCSI._csi && ancestorWithCSI._parent)
-      ancestorWithCSI = ancestorWithCSI._parent;
-
     // Do not report nested async calls to _wrapApiCall.
     isInternal = isInternal || stackTrace.allFrames.filter(f => f.function?.includes('_wrapApiCall')).length > 1;
     if (isInternal)
       delete stackTrace.apiName;
-    const csi = isInternal ? undefined : ancestorWithCSI._csi;
-    const callCookie: { userObject: any } = { userObject: null };
+    const csi = isInternal ? undefined : this._instrumentation;
+    const callCookie: any = {};
 
     try {
       logApiCall(logger, `=> ${apiName} started`, isInternal);

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

@@ -0,0 +1,50 @@
+/**
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ParsedStackTrace } from '../utils/stackTrace';
+
+export interface ClientInstrumentation {
+  addListener(listener: ClientInstrumentationListener): void;
+  removeListener(listener: ClientInstrumentationListener): void;
+  removeAllListeners(): void;
+  onApiCallBegin(apiCall: string, stackTrace: ParsedStackTrace | null, userData: any): void;
+  onApiCallEnd(userData: any, error?: Error): any;
+}
+
+export interface ClientInstrumentationListener {
+  onApiCallBegin?(apiCall: string, stackTrace: ParsedStackTrace | null, userData: any): any;
+  onApiCallEnd?(userData: any, error?: Error): any;
+}
+
+export function createInstrumentation(): ClientInstrumentation {
+  const listeners: ClientInstrumentationListener[] = [];
+  return new Proxy({}, {
+    get: (obj: any, prop: string) => {
+      if (prop === 'addListener')
+        return (listener: ClientInstrumentationListener) => listeners.push(listener);
+      if (prop === 'removeListener')
+        return (listener: ClientInstrumentationListener) => listeners.splice(listeners.indexOf(listener), 1);
+      if (prop === 'removeAllListeners')
+        return () => listeners.splice(0, listeners.length);
+      if (!prop.startsWith('on'))
+        return obj[prop];
+      return async (...params: any[]) => {
+        for (const listener of listeners)
+          await (listener as any)[prop]?.(...params);
+      };
+    },
+  });
+}

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

@@ -62,7 +62,6 @@ export class Connection extends EventEmitter {
   private _rootObject: Root;
   private _closedErrorMessage: string | undefined;
   private _isRemote = false;
-  private _sourceCollector: Set<string> | undefined;
 
   constructor() {
     super();
@@ -89,10 +88,6 @@ export class Connection extends EventEmitter {
     return this._objects.get(guid)!;
   }
 
-  setSourceCollector(collector: Set<string> | undefined) {
-    this._sourceCollector = collector;
-  }
-
   async sendMessageToServer(object: ChannelOwner, method: string, params: any, maybeStackTrace: ParsedStackTrace | null): Promise<any> {
     if (this._closedErrorMessage)
       throw new Error(this._closedErrorMessage);
@@ -100,8 +95,6 @@ export class Connection extends EventEmitter {
     const guid = object._guid;
     const stackTrace: ParsedStackTrace = maybeStackTrace || { frameTexts: [], frames: [], apiName: '', allFrames: [] };
     const { frames, apiName } = stackTrace;
-    if (this._sourceCollector)
-      frames.forEach(f => this._sourceCollector!.add(f.file));
     const id = ++this._lastId;
     const converted = { id, guid, method, params };
     // Do not include metadata in debug logs to avoid noise.

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

@@ -25,18 +25,27 @@ import yazl from 'yazl';
 import { assert, calculateSha1 } from '../utils/utils';
 import { ManualPromise } from '../utils/async';
 import EventEmitter from 'events';
+import { ClientInstrumentationListener } from './clientInstrumentation';
+import { ParsedStackTrace } from '../utils/stackTrace';
 
 export class Tracing implements api.Tracing {
   private _context: BrowserContext;
-  private _sources: Set<string> | undefined;
+  private _sources = new Set<string>();
+  private _instrumentationListener: ClientInstrumentationListener;
 
   constructor(channel: BrowserContext) {
     this._context = channel;
+    this._instrumentationListener = {
+      onApiCallBegin: (apiCall: string, stackTrace: ParsedStackTrace | null) => {
+        for (const frame of stackTrace?.frames || [])
+          this._sources.add(frame.file);
+      }
+    };
   }
 
   async start(options: { name?: string, snapshots?: boolean, screenshots?: boolean, sources?: boolean } = {}) {
-    this._sources = options.sources ? new Set() : undefined;
-    this._context._connection.setSourceCollector(this._sources);
+    if (options.sources)
+      this._context._instrumentation!.addListener(this._instrumentationListener);
     await this._context._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
       await channel.tracingStart(options);
       await channel.tracingStartChunk();
@@ -44,7 +53,7 @@ export class Tracing implements api.Tracing {
   }
 
   async startChunk() {
-    this._context._connection.setSourceCollector(this._sources);
+    this._sources = new Set();
     await this._context._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
       await channel.tracingStartChunk();
     });
@@ -65,7 +74,8 @@ export class Tracing implements api.Tracing {
 
   private async _doStopChunk(channel: channels.BrowserContextChannel, filePath: string | undefined) {
     const sources = this._sources;
-    this._context._connection.setSourceCollector(undefined);
+    this._sources = new Set();
+    this._context._instrumentation!.removeListener(this._instrumentationListener);
     const skipCompress = !this._context._connection.isRemote();
     const result = await channel.tracingStopChunk({ save: !!filePath, skipCompress });
     if (!filePath) {

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

@@ -17,7 +17,6 @@
 
 import * as channels from '../protocol/channels';
 import type { Size } from '../common/types';
-import { ParsedStackTrace } from '../utils/stackTrace';
 export { Size, Point, Rect, Quad, URLMatch, TimeoutOptions, HeadersArray } from '../common/types';
 
 type LoggerSeverity = 'verbose' | 'info' | 'warning' | 'error';
@@ -26,11 +25,6 @@ export interface Logger {
   log(name: string, severity: LoggerSeverity, message: string | Error, args: any[], hints: { color?: string }): void;
 }
 
-export interface ClientSideInstrumentation {
-  onApiCallBegin(apiCall: string, stackTrace: ParsedStackTrace | null): { userObject: any };
-  onApiCallEnd(userData: { userObject: any }, error?: Error): any;
-}
-
 export type StrictOptions = { strict?: boolean };
 export type Headers = { [key: string]: string };
 export type Env = { [key: string]: string | number | boolean | undefined };

packages/playwright-test/src/index.ts

@@ -297,8 +297,8 @@ export const test = _baseTest.extend<TestFixtures, WorkerAndFileFixtures>({
         (context.tracing as any)[kTracingStarted] = false;
         await context.tracing.stop();
       }
-      (context as any)._csi = {
-        onApiCallBegin: (apiCall: string, stackTrace: ParsedStackTrace | null) => {
+      (context as any)._instrumentation.addListener({
+        onApiCallBegin: (apiCall: string, stackTrace: ParsedStackTrace | null, userData: any) => {
           if (apiCall.startsWith('expect.'))
             return { userObject: null };
           const testInfoImpl = testInfo as any;
@@ -309,13 +309,13 @@ export const test = _baseTest.extend<TestFixtures, WorkerAndFileFixtures>({
             canHaveChildren: false,
             forceNoParent: false
           });
-          return { userObject: step };
+          userData.userObject = step;
         },
-        onApiCallEnd: (data: { userObject: any }, error?: Error) => {
-          const step = data.userObject;
+        onApiCallEnd: (userData: any, error?: Error) => {
+          const step = userData.userObject;
           step?.complete(error);
         },
-      };
+      });
     };
 
     const onWillCloseContext = async (context: BrowserContext) => {
@@ -374,7 +374,7 @@ export const test = _baseTest.extend<TestFixtures, WorkerAndFileFixtures>({
     (_browserType as any)._onDidCreateContext = undefined;
     (_browserType as any)._onWillCloseContext = undefined;
     (_browserType as any)._defaultContextOptions = undefined;
-    leftoverContexts.forEach(context => (context as any)._csi = undefined);
+    leftoverContexts.forEach(context => (context as any)._instrumentation.removeAllListeners());
 
     // 5. Collect artifacts from any non-closed contexts.
     await Promise.all(leftoverContexts.map(async context => {

tests/config/browserTest.ts

@@ -147,8 +147,8 @@ export const playwrightFixtures: Fixtures<PlaywrightTestOptions & PlaywrightTest
       context.on('close', () => contexts.get(context).closed = true);
       if (trace)
         await context.tracing.start({ screenshots: true, snapshots: true, sources: true } as any);
-      (context as any)._csi = {
-        onApiCallBegin: (apiCall: string, stackTrace: ParsedStackTrace | null) => {
+      (context as any)._instrumentation.addListener({
+        onApiCallBegin: (apiCall: string, stackTrace: ParsedStackTrace | null, userData: any) => {
           if (apiCall.startsWith('expect.'))
             return { userObject: null };
           const testInfoImpl = testInfo as any;
@@ -159,13 +159,13 @@ export const playwrightFixtures: Fixtures<PlaywrightTestOptions & PlaywrightTest
             canHaveChildren: false,
             forceNoParent: false
           });
-          return { userObject: step };
+          userData.userObject = step;
         },
-        onApiCallEnd: (data: { userObject: any }, error?: Error) => {
-          const step = data.userObject;
+        onApiCallEnd: (userData: any, error?: Error) => {
+          const step = userData.userObject;
           step?.complete(error);
         },
-      };
+      });
       return context;
     });
     await Promise.all([...contexts.keys()].map(async context => {