current zone and run in zone

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

@@ -17,7 +17,8 @@
 import type { EventEmitter } from 'events';
 import { rewriteErrorMessage } from '../utils/stackTrace';
 import { TimeoutError } from './errors';
-import { createGuid } from '../utils';
+import { createGuid, zones } from '../utils';
+import type { Zone } from '../utils';
 import type * as channels from '@protocol/channels';
 import type { ChannelOwner } from './channelOwner';
 
@@ -29,10 +30,13 @@ export class Waiter {
   private _channelOwner: ChannelOwner<channels.EventTargetChannel>;
   private _waitId: string;
   private _error: string | undefined;
+  private _savedZone: Zone | undefined;
 
   constructor(channelOwner: ChannelOwner<channels.EventTargetChannel>, event: string) {
     this._waitId = createGuid();
     this._channelOwner = channelOwner;
+    this._savedZone = zones.currentZone();
+
     this._channelOwner._channel.waitForEventInfo({ info: { waitId: this._waitId, phase: 'before', event } }).catch(() => {});
     this._dispose = [
       () => this._channelOwner._wrapApiCall(async () => {
@@ -46,12 +50,12 @@ export class Waiter {
   }
 
   async waitForEvent<T = void>(emitter: EventEmitter, event: string, predicate?: (arg: T) => boolean | Promise<boolean>): Promise<T> {
-    const { promise, dispose } = waitForEvent(emitter, event, predicate);
+    const { promise, dispose } = waitForEvent(emitter, event, this._savedZone, predicate);
     return await this.waitForPromise(promise, dispose);
   }
 
   rejectOnEvent<T = void>(emitter: EventEmitter, event: string, error: Error | (() => Error), predicate?: (arg: T) => boolean | Promise<boolean>) {
-    const { promise, dispose } = waitForEvent(emitter, event, predicate);
+    const { promise, dispose } = waitForEvent(emitter, event, this._savedZone, predicate);
     this._rejectOn(promise.then(() => { throw (typeof error === 'function' ? error() : error); }), dispose);
   }
 
@@ -103,19 +107,22 @@ export class Waiter {
   }
 }
 
-function waitForEvent<T = void>(emitter: EventEmitter, event: string, predicate?: (arg: T) => boolean | Promise<boolean>): { promise: Promise<T>, dispose: () => void } {
+function waitForEvent<T = void>(emitter: EventEmitter, event: string, savedZone: Zone | undefined, predicate?: (arg: T) => boolean | Promise<boolean>): { promise: Promise<T>, dispose: () => void } {
   let listener: (eventArg: any) => void;
   const promise = new Promise<T>((resolve, reject) => {
     listener = async (eventArg: any) => {
-      try {
-        if (predicate && !(await predicate(eventArg)))
-          return;
-        emitter.removeListener(event, listener);
-        resolve(eventArg);
-      } catch (e) {
-        emitter.removeListener(event, listener);
-        reject(e);
-      }
+      // Reset apiZone and expectZone, but restore step data.
+      await zones.runInZone(savedZone?.copyWithoutTypes(['apiZone', 'expectZone']), async () => {
+        try {
+          if (predicate && !(await predicate(eventArg)))
+            return;
+          emitter.removeListener(event, listener);
+          resolve(eventArg);
+        } catch (e) {
+          emitter.removeListener(event, listener);
+          reject(e);
+        }
+      });
     };
     emitter.addListener(event, listener);
   });

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

@@ -19,49 +19,52 @@ import { AsyncLocalStorage } from 'async_hooks';
 export type ZoneType = 'apiZone' | 'expectZone' | 'stepZone';
 
 class ZoneManager {
-  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone<unknown>|undefined>();
+  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone|undefined>();
 
   run<T, R>(type: ZoneType, data: T, func: () => R): R {
-    const previous = this._asyncLocalStorage.getStore();
-    const zone = new Zone(previous, type, data);
+    const current = this._asyncLocalStorage.getStore();
+    const zone = Zone.createWithData(current, type, data);
+    return this.runInZone(zone, func);
+  }
+
+  runInZone<R>(zone: Zone | undefined, func: () => R): R {
     return this._asyncLocalStorage.run(zone, func);
   }
 
   zoneData<T>(type: ZoneType): T | undefined {
-    for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
-      if (zone.type === type)
-        return zone.data as T;
-    }
-    return undefined;
+    const zone = this._asyncLocalStorage.getStore();
+    return zone?.get(type);
+  }
+
+  currentZone(): Zone | undefined {
+    return this._asyncLocalStorage.getStore();
   }
 
   exitZones<R>(func: () => R): R {
     return this._asyncLocalStorage.run(undefined, func);
   }
-
-  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);
-
-    }
-    // eslint-disable-next-line no-console
-    console.log('zones: ', zones.join(' -> '));
-  }
 }
 
-class Zone<T> {
-  readonly type: ZoneType;
-  readonly data: T;
-  readonly previous: Zone<unknown> | undefined;
+export class Zone {
+  private readonly store: Map<ZoneType, unknown>;
 
-  constructor(previous: Zone<unknown> | undefined, type: ZoneType, data: T) {
-    this.type = type;
-    this.data = data;
-    this.previous = previous;
+  static createWithData(currentZone: Zone | undefined, type: ZoneType, data: unknown) {
+    const store = new Map(currentZone?.store.entries() ?? []);
+    store.set(type, data);
+    return new Zone(store);
+  }
+
+  private constructor(store: Map<ZoneType, unknown>) {
+    this.store = store;
+  }
+
+  copyWithoutTypes(types: ZoneType[]): Zone {
+    const store = new Map(this.store.entries().filter(([type]) => !types.includes(type)));
+    return new Zone(store);
+  }
+
+  get<T>(type: ZoneType): T | undefined {
+    return this.store.get(type) as T | undefined;
   }
 }