chore: replace Zones with AsyncLocalStorage (#30381)

Reference https://github.com/microsoft/playwright/issues/30322

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

@@ -19,7 +19,7 @@ import type * as channels from '@protocol/channels';
 import { maybeFindValidator, ValidationError, type ValidatorContext } from '../protocol/validator';
 import { debugLogger } from '../utils/debugLogger';
 import type { ExpectZone } from '../utils/stackTrace';
-import { captureRawStack, captureLibraryStackTrace, stringifyStackFrames } from '../utils/stackTrace';
+import { captureLibraryStackTrace, stringifyStackFrames } from '../utils/stackTrace';
 import { isUnderTest } from '../utils';
 import { zones } from '../utils/zones';
 import type { ClientInstrumentation } from './clientInstrumentation';
@@ -161,12 +161,11 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
 
   async _wrapApiCall<R>(func: (apiZone: ApiZone) => Promise<R>, isInternal = false): Promise<R> {
     const logger = this._logger;
-    const stack = captureRawStack();
-    const apiZone = zones.zoneData<ApiZone>('apiZone', stack);
+    const apiZone = zones.zoneData<ApiZone>('apiZone');
     if (apiZone)
       return await func(apiZone);
 
-    const stackTrace = captureLibraryStackTrace(stack);
+    const stackTrace = captureLibraryStackTrace();
     let apiName: string | undefined = stackTrace.apiName;
     const frames: channels.StackFrame[] = stackTrace.frames;
 
@@ -175,7 +174,7 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
       apiName = undefined;
 
     // Enclosing zone could have provided the apiName and wallTime.
-    const expectZone = zones.zoneData<ExpectZone>('expectZone', stack);
+    const expectZone = zones.zoneData<ExpectZone>('expectZone');
     const wallTime = expectZone ? expectZone.wallTime : Date.now();
     if (!isInternal && expectZone)
       apiName = expectZone.title;

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

@@ -44,7 +44,7 @@ import { Tracing } from './tracing';
 import { findValidator, ValidationError, type ValidatorContext } from '../protocol/validator';
 import { createInstrumentation } from './clientInstrumentation';
 import type { ClientInstrumentation } from './clientInstrumentation';
-import { formatCallLog, rewriteErrorMessage } from '../utils';
+import { formatCallLog, rewriteErrorMessage, zones } from '../utils';
 
 class Root extends ChannelOwner<channels.RootChannel> {
   constructor(connection: Connection) {
@@ -136,7 +136,9 @@ export class Connection extends EventEmitter {
     const metadata: channels.Metadata = { wallTime, apiName, location, internal: !apiName };
     if (this._tracingCount && frames && type !== 'LocalUtils')
       this._localUtils?._channel.addStackToTracingNoReply({ callData: { stack: frames, id } }).catch(() => {});
-    this.onmessage({ ...message, metadata });
+    // We need to exit zones before calling into the server, otherwise
+    // when we receive events from the server, we would be in an API zone.
+    zones.exitZones(() => this.onmessage({ ...message, metadata }));
     return await new Promise((resolve, reject) => this._callbacks.set(id, { resolve, reject, apiName, type, method }));
   }
 

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

@@ -48,8 +48,8 @@ export function captureRawStack(): RawStack {
   return stack.split('\n');
 }
 
-export function captureLibraryStackTrace(rawStack?: RawStack): { frames: StackFrame[], apiName: string } {
-  const stack = rawStack || captureRawStack();
+export function captureLibraryStackTrace(): { frames: StackFrame[], apiName: string } {
+  const stack = captureRawStack();
 
   const isTesting = isUnderTest();
   type ParsedFrame = {

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

@@ -14,88 +14,53 @@
  * limitations under the License.
  */
 
-import type { RawStack } from './stackTrace';
-import { captureRawStack } from './stackTrace';
+import { AsyncLocalStorage } from 'async_hooks';
 
 export type ZoneType = 'apiZone' | 'expectZone' | 'stepZone';
 
 class ZoneManager {
-  lastZoneId = 0;
-  readonly _zones = new Map<number, Zone<any>>();
+  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone<unknown>|undefined>();
 
   run<T, R>(type: ZoneType, data: T, func: (data: T) => R): R {
-    return new Zone<T>(this, ++this.lastZoneId, type, data).run(func);
+    const previous = this._asyncLocalStorage.getStore();
+    const zone = new Zone(previous, type, data);
+    return this._asyncLocalStorage.run(zone, () => func(data));
   }
 
-  zoneData<T>(type: ZoneType, rawStack: RawStack): T | null {
-    for (const line of rawStack) {
-      for (const zoneId of zoneIds(line)) {
-        const zone = this._zones.get(zoneId);
-        if (zone && zone.type === type)
-          return zone.data;
-      }
+  zoneData<T>(type: ZoneType): T | null {
+    for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
+      if (zone.type === type)
+        return zone.data as T;
     }
     return null;
   }
 
-  preserve<T>(callback: () => Promise<T>): Promise<T> {
-    const rawStack = captureRawStack();
-    const refs: number[] = [];
-    for (const line of rawStack)
-      refs.push(...zoneIds(line));
-    Object.defineProperty(callback, 'name', { value: `__PWZONE__[${refs.join(',')}]-refs` });
-    return callback();
+  exitZones<R>(func: () => R): R {
+    return this._asyncLocalStorage.run(undefined, func);
   }
-}
 
-function zoneIds(line: string): number[] {
-  const index = line.indexOf('__PWZONE__[');
-  if (index === -1)
-    return [];
-  return line.substring(index + '__PWZONE__['.length, line.indexOf(']', index)).split(',').map(s => +s);
+  printZones() {
+    const zones = [];
+    for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
+      let str = zone.type;
+      if (zone.type === 'apiZone')
+        str += `(${(zone.data as any).apiName})`;
+      zones.push(str);
+      
+    }
+    console.log('zones: ', zones.join(' -> '));
+  }
 }
 
 class Zone<T> {
-  private _manager: ZoneManager;
-  readonly id: number;
   readonly type: ZoneType;
-  data: T;
-  readonly wallTime: number;
+  readonly data: T;
+  readonly previous: Zone<unknown> | undefined;
 
-  constructor(manager: ZoneManager, id: number, type: ZoneType, data: T) {
-    this._manager = manager;
-    this.id = id;
+  constructor(previous: Zone<unknown> | undefined, type: ZoneType, data: T) {
     this.type = type;
     this.data = data;
-    this.wallTime = Date.now();
-  }
-
-  run<R>(func: (data: T) => R): R {
-    this._manager._zones.set(this.id, this);
-    Object.defineProperty(func, 'name', { value: `__PWZONE__[${this.id}]-${this.type}` });
-    return runWithFinally(() => func(this.data), () => {
-      this._manager._zones.delete(this.id);
-    });
-  }
-}
-
-export function runWithFinally<R>(func: () => R, finallyFunc: Function): R {
-  try {
-    const result = func();
-    if (result instanceof Promise) {
-      return result.then(r => {
-        finallyFunc();
-        return r;
-      }).catch(e => {
-        finallyFunc();
-        throw e;
-      }) as any;
-    }
-    finallyFunc();
-    return result;
-  } catch (e) {
-    finallyFunc();
-    throw e;
+    this.previous = previous;
   }
 }
 

packages/playwright/src/matchers/expect.ts

@@ -257,8 +257,7 @@ class ExpectMetaInfoProxyHandler implements ProxyHandler<any> {
 
       // This looks like it is unnecessary, but it isn't - we need to filter
       // out all the frames that belong to the test runner from caught runtime errors.
-      const rawStack = captureRawStack();
-      const stackFrames = filteredStackTrace(rawStack);
+      const stackFrames = filteredStackTrace(captureRawStack());
 
       // Enclose toPass in a step to maintain async stacks, toPass matcher is always async.
       const stepInfo = {
@@ -287,7 +286,7 @@ class ExpectMetaInfoProxyHandler implements ProxyHandler<any> {
       try {
         const expectZone: ExpectZone | null = matcherName !== 'toPass' ? { title, wallTime } : null;
         const callback = () => matcher.call(target, ...args);
-        const result = expectZone ? zones.run<ExpectZone, any>('expectZone', expectZone, callback) : zones.preserve(callback);
+        const result = expectZone ? zones.run<ExpectZone, any>('expectZone', expectZone, callback) : callback();
         if (result instanceof Promise)
           return result.then(finalizer).catch(reportStepError);
         finalizer();

packages/playwright/src/matchers/toMatchSnapshot.ts

@@ -18,7 +18,7 @@ import type { Locator, Page } from 'playwright-core';
 import type { ExpectScreenshotOptions, Page as PageEx } from 'playwright-core/lib/client/page';
 import { currentTestInfo, currentExpectTimeout } from '../common/globals';
 import type { ImageComparatorOptions, Comparator } from 'playwright-core/lib/utils';
-import { getComparator, sanitizeForFilePath, zones } from 'playwright-core/lib/utils';
+import { getComparator, sanitizeForFilePath } from 'playwright-core/lib/utils';
 import {
   addSuffixToFilePath,
   trimLongString, callLogText,
@@ -367,19 +367,7 @@ export async function toHaveScreenshot(
   if (!helper.snapshotPath.toLowerCase().endsWith('.png'))
     throw new Error(`Screenshot name "${path.basename(helper.snapshotPath)}" must have '.png' extension`);
   expectTypes(pageOrLocator, ['Page', 'Locator'], 'toHaveScreenshot');
-  return await zones.preserve(async () => {
-    // Loading from filesystem resets zones.
-    const style = await loadScreenshotStyles(helper.options.stylePath);
-    return toHaveScreenshotContinuation.call(this, helper, page, locator, style);
-  });
-}
-
-async function toHaveScreenshotContinuation(
-  this: ExpectMatcherContext,
-  helper: SnapshotHelper,
-  page: PageEx,
-  locator: Locator | undefined,
-  style?: string) {
+  const style = await loadScreenshotStyles(helper.options.stylePath);
   const expectScreenshotOptions: ExpectScreenshotOptions = {
     locator,
     animations: helper.options.animations ?? 'disabled',

packages/playwright/src/worker/testInfo.ts

@@ -246,14 +246,13 @@ export class TestInfoImpl implements TestInfo {
 
   _addStep(data: Omit<TestStepInternal, 'complete' | 'stepId' | 'steps'>): TestStepInternal {
     const stepId = `${data.category}@${++this._lastStepId}`;
-    const rawStack = captureRawStack();
 
     let parentStep: TestStepInternal | undefined;
     if (data.isStage) {
       // Predefined stages form a fixed hierarchy - use the current one as parent.
       parentStep = this._findLastStageStep();
     } else {
-      parentStep = zones.zoneData<TestStepInternal>('stepZone', rawStack!) || undefined;
+      parentStep = zones.zoneData<TestStepInternal>('stepZone') || undefined;
       if (!parentStep && data.category !== 'test.step') {
         // API steps (but not test.step calls) can be nested by time, instead of by stack.
         // However, do not nest chains of route.continue by checking the title.
@@ -265,7 +264,7 @@ export class TestInfoImpl implements TestInfo {
       }
     }
 
-    const filteredStack = filteredStackTrace(rawStack);
+    const filteredStack = filteredStackTrace(captureRawStack());
     data.boxedStack = parentStep?.boxedStack;
     if (!data.boxedStack && data.box) {
       data.boxedStack = filteredStack.slice(1);

tests/playwright-test/test-step.spec.ts

@@ -52,6 +52,15 @@ class Reporter {
     this.suite = suite;
   }
 
+  // For easier debugging.
+  onStdOut(data) {
+    process.stdout.write(data.toString());
+  }
+  // For easier debugging.
+  onStdErr(data) {
+    process.stderr.write(data.toString());
+  }
+
   printStep(step, indent) {
     let location = '';
     if (step.location)
@@ -867,7 +876,6 @@ test('step inside expect.toPass', async ({ runInlineTest }) => {
   }, { reporter: '', workers: 1 });
 
   expect(result.exitCode).toBe(0);
-  console.log(result.output);
   expect(stripAnsi(result.output)).toBe(`
 hook      |Before Hooks
 test.step |step 1 @ a.test.ts:4