chore: wrap expect call in a zone (#21191)

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

@@ -18,7 +18,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 type { ParsedStackTrace } from '../utils/stackTrace';
+import type { ExpectZone, ParsedStackTrace } from '../utils/stackTrace';
 import { captureRawStack, captureLibraryStackTrace } from '../utils/stackTrace';
 import { isUnderTest } from '../utils';
 import { zones } from '../utils/zones';
@@ -138,9 +138,8 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
           if (validator) {
             return (params: any) => {
               return this._wrapApiCall(apiZone => {
-                const { stackTrace, csi, callCookie } = apiZone.reported ? { csi: undefined, callCookie: undefined, stackTrace: null } : apiZone;
+                const { stackTrace, csi, callCookie, wallTime } = apiZone.reported ? { csi: undefined, callCookie: undefined, stackTrace: null, wallTime: undefined } : apiZone;
                 apiZone.reported = true;
-                const wallTime = Date.now();
                 if (csi && stackTrace && stackTrace.apiName)
                   csi.onApiCallBegin(renderCallWithParams(stackTrace.apiName, params), stackTrace, wallTime, callCookie);
                 return this._connection.sendMessageToServer(this, this._type, prop, validator(params, '', { tChannelImpl: tChannelImplToWire, binary: this._connection.isRemote() ? 'toBase64' : 'buffer' }), stackTrace, wallTime);
@@ -165,14 +164,20 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
     const stackTrace = captureLibraryStackTrace(stack);
     if (isInternal)
       delete stackTrace.apiName;
+
+    // Enclosing zone could have provided the apiName and wallTime.
+    const expectZone = zones.zoneData<ExpectZone>('expectZone', stack);
+    const wallTime = expectZone ? expectZone.wallTime : Date.now();
+    if (!isInternal && expectZone)
+      stackTrace.apiName = expectZone.title;
+
     const csi = isInternal ? undefined : this._instrumentation;
     const callCookie: any = {};
 
     const { apiName, frameTexts } = stackTrace;
-
     try {
       logApiCall(logger, `=> ${apiName} started`, isInternal);
-      const apiZone = { stackTrace, isInternal, reported: false, csi, callCookie };
+      const apiZone = { stackTrace, isInternal, reported: false, csi, callCookie, wallTime };
       const result = await zones.run<ApiZone, R>('apiZone', apiZone, async () => {
         return await func(apiZone);
       });
@@ -243,4 +248,5 @@ type ApiZone = {
   reported: boolean;
   csi: ClientInstrumentation | undefined;
   callCookie: any;
+  wallTime: number;
 };

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

@@ -15,7 +15,6 @@
  */
 
 import type { ParsedStackTrace } from '../utils/stackTrace';
-
 export interface ClientInstrumentation {
   addListener(listener: ClientInstrumentationListener): void;
   removeListener(listener: ClientInstrumentationListener): void;

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

@@ -111,7 +111,7 @@ export class Connection extends EventEmitter {
     this._stackCollectors.delete(collector);
   }
 
-  async sendMessageToServer(object: ChannelOwner, type: string, method: string, params: any, stackTrace: ParsedStackTrace | null, wallTime: number): Promise<any> {
+  async sendMessageToServer(object: ChannelOwner, type: string, method: string, params: any, stackTrace: ParsedStackTrace | null, wallTime: number | undefined): Promise<any> {
     if (this._closedErrorMessage)
       throw new Error(this._closedErrorMessage);
 

packages/playwright-core/src/utils/index.ts

@@ -38,3 +38,4 @@ export * from './timeoutRunner';
 export * from './traceUtils';
 export * from './userAgent';
 export * from './zipFile';
+export * from './zones';

packages/playwright-core/src/utils/stackTrace.ts

@@ -49,16 +49,18 @@ export type ParsedStackTrace = {
   apiName: string | undefined;
 };
 
-export function captureRawStack(): string {
+export type RawStack = string[];
+
+export function captureRawStack(): RawStack {
   const stackTraceLimit = Error.stackTraceLimit;
   Error.stackTraceLimit = 30;
   const error = new Error();
-  const stack = error.stack!;
+  const stack = error.stack || '';
   Error.stackTraceLimit = stackTraceLimit;
-  return stack;
+  return stack.split('\n');
 }
 
-export function captureLibraryStackTrace(rawStack?: string): ParsedStackTrace {
+export function captureLibraryStackTrace(rawStack?: RawStack): ParsedStackTrace {
   const stack = rawStack || captureRawStack();
 
   const isTesting = isUnderTest();
@@ -67,7 +69,7 @@ export function captureLibraryStackTrace(rawStack?: string): ParsedStackTrace {
     frameText: string;
     isPlaywrightLibrary: boolean;
   };
-  let parsedFrames = stack.split('\n').map(line => {
+  let parsedFrames = stack.map(line => {
     const frame = parseStackTraceLine(line);
     if (!frame || !frame.fileName)
       return null;
@@ -90,16 +92,7 @@ export function captureLibraryStackTrace(rawStack?: string): ParsedStackTrace {
   let apiName = '';
   const allFrames = parsedFrames;
 
-  // Use stack trap for the API annotation, if available.
-  for (let i = parsedFrames.length - 1; i >= 0; i--) {
-    const parsedFrame = parsedFrames[i];
-    if (parsedFrame.frame.function?.startsWith('__PWTRAP__[')) {
-      apiName = parsedFrame.frame.function!.substring('__PWTRAP__['.length, parsedFrame.frame.function!.length - 1);
-      break;
-    }
-  }
-
-  // Otherwise, deepest transition between non-client code calling into client
+  // Deepest transition between non-client code calling into client
   // code is the api entry.
   for (let i = 0; i < parsedFrames.length - 1; i++) {
     const parsedFrame = parsedFrames[i];
@@ -142,3 +135,8 @@ export function splitErrorMessage(message: string): { name: string, message: str
     message: separationIdx !== -1 && separationIdx + 2 <= message.length ? message.substring(separationIdx + 2) : message,
   };
 }
+
+export type ExpectZone = {
+  title: string;
+  wallTime: number;
+};

packages/playwright-core/src/utils/zones.ts

@@ -14,25 +14,20 @@
  * limitations under the License.
  */
 
-import { captureRawStack } from './stackTrace';
+import type { RawStack } from './stackTrace';
+
+export type ZoneType = 'apiZone' | 'expectZone';
 
 class ZoneManager {
   lastZoneId = 0;
-  readonly _zones = new Map<number, Zone>();
+  readonly _zones = new Map<number, Zone<any>>();
 
-  constructor() {
+  run<T, R>(type: ZoneType, data: T, func: (data: T) => R | Promise<R>): R | Promise<R> {
+    return new Zone<T>(this, ++this.lastZoneId, type, data).run(func);
   }
 
-  async run<T, R>(type: string, data: T, func: () => Promise<R>): Promise<R> {
-    const zone = new Zone(this, ++this.lastZoneId, type, data);
-    this._zones.set(zone.id, zone);
-    return zone.run(func);
-  }
-
-  zoneData<T>(type: string, rawStack?: string): T | null {
-    const stack = rawStack || captureRawStack();
-
-    for (const line of stack.split('\n')) {
+  zoneData<T>(type: ZoneType, rawStack: RawStack): T | null {
+    for (const line of rawStack) {
       const index = line.indexOf('__PWZONE__[');
       if (index !== -1) {
         const zoneId = + line.substring(index + '__PWZONE__['.length, line.indexOf(']', index));
@@ -45,27 +40,48 @@ class ZoneManager {
   }
 }
 
-class Zone {
+class Zone<T> {
   private _manager: ZoneManager;
   readonly id: number;
-  readonly type: string;
-  readonly data: any = {};
+  readonly type: ZoneType;
+  data: T;
+  readonly wallTime: number;
 
-  constructor(manager: ZoneManager, id: number, type: string, data: any) {
+  constructor(manager: ZoneManager, id: number, type: ZoneType, data: T) {
     this._manager = manager;
     this.id = id;
     this.type = type;
     this.data = data;
+    this.wallTime = Date.now();
   }
 
-  async run<R>(func: () => Promise<R>): Promise<R> {
+  run<R>(func: (data: T) => R | Promise<R>): R | Promise<R> {
+    this._manager._zones.set(this.id, this);
     Object.defineProperty(func, 'name', { value: `__PWZONE__[${this.id}]` });
-    try {
-      return await func();
-    } finally {
+    return runWithFinally(() => func(this.data), () => {
       this._manager._zones.delete(this.id);
+    });
   }
 }
+
+export function runWithFinally<R>(func: () => R | Promise<R>, finallyFunc: Function): R | Promise<R> {
+  try {
+    const result = func();
+    if (result instanceof Promise) {
+      return result.then(r => {
+        finallyFunc();
+        return r;
+      }).catch(e => {
+        finallyFunc();
+        throw e;
+      });
+    }
+    finallyFunc();
+    return result;
+  } catch (e) {
+    finallyFunc();
+    throw e;
+  }
 }
 
 export const zones = new ZoneManager();

packages/playwright-test/src/matchers/expect.ts

@@ -14,7 +14,8 @@
  * limitations under the License.
  */
 
-import { pollAgainstTimeout } from 'playwright-core/lib/utils';
+import { captureRawStack, pollAgainstTimeout } from 'playwright-core/lib/utils';
+import type { ExpectZone } from 'playwright-core/lib/utils';
 import path from 'path';
 import {
   toBeChecked,
@@ -51,6 +52,7 @@ import {
   RECEIVED_COLOR,
   printReceived,
 } from '../common/expectBundle';
+import { zones } from 'playwright-core/lib/utils';
 
 // from expect/build/types
 export type SyncExpectationResult = {
@@ -196,17 +198,19 @@ class ExpectMetaInfoProxyHandler {
       if (!testInfo)
         return matcher.call(target, ...args);
 
-      const stackFrames = filteredStackTrace(new Error());
+      const rawStack = captureRawStack();
+      const stackFrames = filteredStackTrace(rawStack);
       const frame = stackFrames[0];
       const customMessage = this._info.message || '';
       const defaultTitle = `expect${this._info.isPoll ? '.poll' : ''}${this._info.isSoft ? '.soft' : ''}${this._info.isNot ? '.not' : ''}.${matcherName}`;
+      const wallTime = Date.now();
       const step = testInfo._addStep({
         location: frame && frame.fileName ? { file: path.resolve(process.cwd(), frame.fileName), line: frame.line || 0, column: frame.column || 0 } : undefined,
         category: 'expect',
         title: trimLongString(customMessage || defaultTitle, 1024),
         canHaveChildren: true,
         forceNoParent: false,
-        wallTime: Date.now()
+        wallTime
       });
       testInfo.currentStep = step;
 
@@ -242,7 +246,8 @@ class ExpectMetaInfoProxyHandler {
       };
 
       try {
-        const result = namedFunction(defaultTitle)(() => {
+        const expectZone: ExpectZone = { title: defaultTitle, wallTime };
+        const result = zones.run<ExpectZone, any>('expectZone', expectZone, () => {
           return matcher.call(target, ...args);
         });
         if (result instanceof Promise)
@@ -256,14 +261,6 @@ class ExpectMetaInfoProxyHandler {
   }
 }
 
-function namedFunction(name: string) {
-  const result = function(callback: any) {
-    return callback();
-  };
-  Object.defineProperty(result, 'name', { value: '__PWTRAP__[' + name + ']' });
-  return result;
-}
-
 async function pollMatcher(matcherName: any, isNot: boolean, pollIntervals: number[] | undefined, timeout: number, generator: () => any, ...args: any[]) {
   const result = await pollAgainstTimeout<Error|undefined>(async () => {
     const value = await generator();

packages/playwright-test/src/util.ts

@@ -23,6 +23,7 @@ import url from 'url';
 import { colors, debug, minimatch, parseStackTraceLine } from 'playwright-core/lib/utilsBundle';
 import type { TestInfoError, Location } from './common/types';
 import { calculateSha1, isRegExp, isString } from 'playwright-core/lib/utils';
+import type { RawStack } from 'playwright-core/lib/utils';
 
 const PLAYWRIGHT_TEST_PATH = path.join(__dirname, '..');
 const PLAYWRIGHT_CORE_PATH = path.dirname(require.resolve('playwright-core/package.json'));
@@ -30,15 +31,15 @@ const PLAYWRIGHT_CORE_PATH = path.dirname(require.resolve('playwright-core/packa
 export function filterStackTrace(e: Error) {
   if (process.env.PWDEBUGIMPL)
     return;
-  const stackLines = stringifyStackFrames(filteredStackTrace(e));
+  const stackLines = stringifyStackFrames(filteredStackTrace(e.stack?.split('\n') || []));
   const message = e.message;
   e.stack = `${e.name}: ${e.message}\n${stackLines.join('\n')}`;
   e.message = message;
 }
 
-export function filteredStackTrace(e: Error): StackFrameData[] {
+export function filteredStackTrace(rawStack: RawStack): StackFrameData[] {
   const frames: StackFrameData[] = [];
-  for (const line of e.stack?.split('\n') || []) {
+  for (const line of rawStack) {
     const frame = parseStackTraceLine(line);
     if (!frame || !frame.fileName)
       continue;