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)
-      callback(error);
+    for (const promise of this._raceAgainstEvaluationStallingEventsPromises)
+      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 context = this._existingMainContext();
-    if (!context)
-      throw new Error('Frame does not yet have a main execution context');
 
-    let callback = () => {};
-    const frameInvalidated = new Promise<void>((f, r) => callback = r);
-    this._nonStallingEvaluations.add(callback);
+    const promise = new ManualPromise<T>();
+    this._raceAgainstEvaluationStallingEventsPromises.add(promise);
     try {
       return await Promise.race([
-        context.rawEvaluateJSON(expression),
-        frameInvalidated
+        cb(),
+        promise
       ]);
     } finally {
-      this._nonStallingEvaluations.delete(callback);
+      this._raceAgainstEvaluationStallingEventsPromises.delete(promise);
     }
   }
 
-  async nonStallingEvaluateInExistingContext(expression: string, isFunction: boolean|undefined, world: types.World): Promise<any> {
-    if (this._pendingDocument)
-      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');
+  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.evaluateExpression(expression, isFunction),
-        frameInvalidated
-      ]);
-    } finally {
-      this._nonStallingEvaluations.delete(callback);
-    }
+  nonStallingEvaluateInExistingContext(expression: string, isFunction: boolean|undefined, world: types.World): Promise<any> {
+    return this.raceAgainstEvaluationStallingEvents(() => {
+      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);
+    });
   }
 
   private _recalculateLifecycle() {
@@ -1168,10 +1162,12 @@ export class Frame extends SdkObject {
   }
 
   async hideHighlight() {
-    const context = await this._utilityContext();
-    const injectedScript = await context.injectedScript();
-    return await injectedScript.evaluate(injected => {
-      return injected.hideHighlight();
+    return this.raceAgainstEvaluationStallingEvents(async () => {
+      const context = await this._utilityContext();
+      const injectedScript = await context.injectedScript();
+      return await injectedScript.evaluate(injected => {
+        return injected.hideHighlight();
+      });
     });
   }
 

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,13 +265,14 @@ 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.
 
-    await this._page.hideHighlight();
+    if (hasHighlight)
+      await this._page.hideHighlight();
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 
     if (shouldSetDefaultBackground)

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;
+    });
   });
 });