fix(screenshot): do not stall on hideHighlight (#12764)

packages/playwright-core/src/server/frames.ts

@@ -441,7 +441,7 @@ export class Frame extends SdkObject {
   private _setContentCounter = 0;
   readonly _detachedPromise: Promise<void>;
   private _detachedCallback = () => {};
-  private _nonStallingEvaluations = new Set<(error: Error) => void>();
+  private _raceAgainstEvaluationStallingEventsPromises = new Set<ManualPromise<any>>();
 
   constructor(page: Page, id: string, parentFrame: Frame | null) {
     super(page, 'frame');
@@ -500,53 +500,47 @@ export class Frame extends SdkObject {
   }
 
   _invalidateNonStallingEvaluations(message: string) {
-    if (!this._nonStallingEvaluations)
+    if (!this._raceAgainstEvaluationStallingEventsPromises.size)
       return;
     const error = new Error(message);
-    for (const callback of this._nonStallingEvaluations)
+    for (const promise of this._raceAgainstEvaluationStallingEventsPromises)
-      callback(error);
+      promise.reject(error);
   }
 
-  async nonStallingRawEvaluateInExistingMainContext(expression: string): Promise<any> {
+  async raceAgainstEvaluationStallingEvents<T>(cb: () => Promise<T>): Promise<T> {
     if (this._pendingDocument)
       throw new Error('Frame is currently attempting a navigation');
     if (this._page._frameManager._openedDialogs.size)
       throw new Error('Open JavaScript dialog prevents evaluation');
+
+    const promise = new ManualPromise<T>();
+    this._raceAgainstEvaluationStallingEventsPromises.add(promise);
+    try {
+      return await Promise.race([
+        cb(),
+        promise
+      ]);
+    } finally {
+      this._raceAgainstEvaluationStallingEventsPromises.delete(promise);
+    }
+  }
+
+  nonStallingRawEvaluateInExistingMainContext(expression: string): Promise<any> {
+    return this.raceAgainstEvaluationStallingEvents(() => {
       const context = this._existingMainContext();
       if (!context)
         throw new Error('Frame does not yet have a main execution context');
-
+      return context.rawEvaluateJSON(expression);
-    let callback = () => {};
+    });
-    const frameInvalidated = new Promise<void>((f, r) => callback = r);
-    this._nonStallingEvaluations.add(callback);
-    try {
-      return await Promise.race([
-        context.rawEvaluateJSON(expression),
-        frameInvalidated
-      ]);
-    } finally {
-      this._nonStallingEvaluations.delete(callback);
-    }
   }
 
-  async nonStallingEvaluateInExistingContext(expression: string, isFunction: boolean|undefined, world: types.World): Promise<any> {
+  nonStallingEvaluateInExistingContext(expression: string, isFunction: boolean|undefined, world: types.World): Promise<any> {
-    if (this._pendingDocument)
+    return this.raceAgainstEvaluationStallingEvents(() => {
-      throw new Error('Frame is currently attempting a navigation');
       const context = this._contextData.get(world)?.context;
       if (!context)
         throw new Error('Frame does not yet have the execution context');
-
+      return context.evaluateExpression(expression, isFunction);
-    let callback = () => {};
+    });
-    const frameInvalidated = new Promise<void>((f, r) => callback = r);
-    this._nonStallingEvaluations.add(callback);
-    try {
-      return await Promise.race([
-        context.evaluateExpression(expression, isFunction),
-        frameInvalidated
-      ]);
-    } finally {
-      this._nonStallingEvaluations.delete(callback);
-    }
   }
 
   private _recalculateLifecycle() {
@@ -1168,11 +1162,13 @@ export class Frame extends SdkObject {
   }
 
   async hideHighlight() {
+    return this.raceAgainstEvaluationStallingEvents(async () => {
       const context = await this._utilityContext();
       const injectedScript = await context.injectedScript();
       return await injectedScript.evaluate(injected => {
         return injected.hideHighlight();
       });
+    });
   }
 
   private async _elementState(metadata: CallMetadata, selector: string, state: ElementStateWithoutStable, options: types.QueryOnSelectorOptions = {}): Promise<boolean> {

packages/playwright-core/src/server/screenshotter.ts

@@ -236,6 +236,9 @@ export class Screenshotter {
   }
 
   async _maskElements(progress: Progress, options: ScreenshotOptions) {
+    if (!options.mask || !options.mask.length)
+      return false;
+
     const framesToParsedSelectors: MultiMap<Frame, ParsedSelector> = new MultiMap();
     await Promise.all((options.mask || []).map(async ({ frame, selector }) => {
       const pair = await frame.resolveFrameForSelectorNoWait(selector);
@@ -248,6 +251,7 @@ export class Screenshotter {
       await frame.maskSelectors(framesToParsedSelectors.get(frame));
     }));
     progress.cleanupWhenAborted(() => this._page.hideHighlight());
+    return true;
   }
 
   private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, fitsViewport: boolean, options: ScreenshotOptions): Promise<Buffer> {
@@ -261,12 +265,13 @@ export class Screenshotter {
     }
     progress.throwIfAborted(); // Avoid extra work.
 
-    await this._maskElements(progress, options);
+    const hasHighlight = await this._maskElements(progress, options);
     progress.throwIfAborted(); // Avoid extra work.
 
     const buffer = await this._page._delegate.takeScreenshot(progress, format, documentRect, viewportRect, options.quality, fitsViewport, options.size || 'device');
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 
+    if (hasHighlight)
       await this._page.hideHighlight();
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 

tests/page/page-screenshot.spec.ts

@@ -17,6 +17,7 @@
 
 import { test as it, expect } from './pageTest';
 import { verifyViewport, attachFrame } from '../config/utils';
+import type { Route } from 'playwright-core';
 import path from 'path';
 import fs from 'fs';
 import os from 'os';
@@ -429,6 +430,20 @@ it.describe('page screenshot', () => {
       const screenshot2 = await page.screenshot();
       expect(screenshot1.equals(screenshot2)).toBe(true);
     });
+
+    it('should work when subframe has stalled navigation', async ({ page, server }) => {
+      let cb;
+      const routeReady = new Promise<Route>(f => cb = f);
+      await page.route('**/subframe.html', cb); // Stalling subframe.
+
+      await page.goto(server.EMPTY_PAGE);
+      const done = page.setContent(`<iframe src='/subframe.html'></iframe>`);
+      const route = await routeReady;
+
+      await page.screenshot({ mask: [ page.locator('non-existent') ] });
+      await route.fulfill({ body: '' });
+      await done;
+    });
   });
 });