chore: make screenshot progress "strict"

packages/playwright-core/src/server/dispatchers/frameDispatcher.ts

@@ -246,7 +246,7 @@ export class FrameDispatcher extends Dispatcher<Frame, channels.FrameChannel, Br
   }
 
   async waitForFunction(params: channels.FrameWaitForFunctionParams, metadata: CallMetadata): Promise<channels.FrameWaitForFunctionResult> {
-    return { handle: ElementHandleDispatcher.fromJSOrElementHandle(this, await this._frame._waitForFunctionExpression(metadata, params.expression, params.isFunction, parseArgument(params.arg), params)) };
+    return { handle: ElementHandleDispatcher.fromJSOrElementHandle(this, await this._frame.waitForFunctionExpression(metadata, params.expression, params.isFunction, parseArgument(params.arg), params)) };
   }
 
   async title(params: channels.FrameTitleParams, metadata: CallMetadata): Promise<channels.FrameTitleResult> {

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

@@ -839,7 +839,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
   }
 
   async screenshot(metadata: CallMetadata, options: ScreenshotOptions & types.TimeoutOptions): Promise<Buffer> {
-    const controller = new ProgressController(metadata, this);
+    const controller = new ProgressController(metadata, this, 'strict');
     return controller.run(
         progress => this._page.screenshotter.screenshotElement(progress, this, options),
         options.timeout);

packages/playwright-core/src/server/firefox/ffPage.ts

@@ -419,7 +419,6 @@ export class FFPage implements PageDelegate {
         height: viewportRect!.height,
       };
     }
-    progress.throwIfAborted();
     const { data } = await this._session.send('Page.screenshot', {
       mimeType: ('image/' + format) as ('image/png' | 'image/jpeg'),
       clip: documentRect,

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

@@ -1135,7 +1135,7 @@ export class Frame extends SdkObject {
 
   async rafrafTimeoutScreenshotElementWithProgress(progress: Progress, selector: string, timeout: number, options: ScreenshotOptions): Promise<Buffer> {
     return await this._retryWithProgressIfNotConnected(progress, selector, true /* strict */, true /* performActionPreChecks */, async handle => {
-      await handle._frame.rafrafTimeout(timeout);
+      await handle._frame.rafrafTimeout(progress, timeout);
       return await this._page.screenshotter.screenshotElement(progress, handle, options);
     });
   }
@@ -1482,63 +1482,67 @@ export class Frame extends SdkObject {
     return { matches, received };
   }
 
-  async _waitForFunctionExpression<R>(metadata: CallMetadata, expression: string, isFunction: boolean | undefined, arg: any, options: types.WaitForFunctionOptions, world: types.World = 'main'): Promise<js.SmartHandle<R>> {
-    const controller = new ProgressController(metadata, this);
+  async waitForFunctionExpression<R>(metadata: CallMetadata, expression: string, isFunction: boolean | undefined, arg: any, options: types.WaitForFunctionOptions): Promise<js.SmartHandle<R>> {
+    const controller = new ProgressController(metadata, this, 'strict');
+    return controller.run(progress => this.waitForFunctionExpressionImpl(progress, expression, isFunction, arg, options, 'main'), options.timeout);
+  }
+
+  async waitForFunctionExpressionImpl<R>(progress: Progress, expression: string, isFunction: boolean | undefined, arg: any, options: { pollingInterval?: number }, world: types.World = 'main'): Promise<js.SmartHandle<R>> {
     if (typeof options.pollingInterval === 'number')
       assert(options.pollingInterval > 0, 'Cannot poll with non-positive interval: ' + options.pollingInterval);
     expression = js.normalizeEvaluationExpression(expression, isFunction);
-    return controller.run(async progress => {
-      return this.retryWithProgressAndTimeouts(progress, [100], async () => {
-        const context = world === 'main' ? await this._mainContext() : await this._utilityContext();
-        const injectedScript = await context.injectedScript();
-        const handle = await injectedScript.evaluateHandle((injected, { expression, isFunction, polling, arg }) => {
-          const predicate = (): R => {
-            // NOTE: make sure to use `globalThis.eval` instead of `self.eval` due to a bug with sandbox isolation
-            // in firefox.
-            // See https://bugzilla.mozilla.org/show_bug.cgi?id=1814898
-            let result = globalThis.eval(expression);
-            if (isFunction === true) {
+    return this.retryWithProgressAndTimeouts(progress, [100], async () => {
+      const context = world === 'main' ? await progress.race(this._mainContext()) : await progress.race(this._utilityContext());
+      const injectedScript = await progress.race(context.injectedScript());
+      const handle = await progress.race(injectedScript.evaluateHandle((injected, { expression, isFunction, polling, arg }) => {
+        const predicate = (): R => {
+          // NOTE: make sure to use `globalThis.eval` instead of `self.eval` due to a bug with sandbox isolation
+          // in firefox.
+          // See https://bugzilla.mozilla.org/show_bug.cgi?id=1814898
+          let result = globalThis.eval(expression);
+          if (isFunction === true) {
+            result = result(arg);
+          } else if (isFunction === false) {
+            result = result;
+          } else {
+            // auto detect.
+            if (typeof result === 'function')
               result = result(arg);
-            } else if (isFunction === false) {
-              result = result;
-            } else {
-              // auto detect.
-              if (typeof result === 'function')
-                result = result(arg);
-            }
-            return result;
-          };
-
-          let fulfill: (result: R) => void;
-          let reject: (error: Error) => void;
-          let aborted = false;
-          const result = new Promise<R>((f, r) => { fulfill = f; reject = r; });
-
-          const next = () => {
-            if (aborted)
+          }
+          return result;
+        };
+
+        let fulfill: (result: R) => void;
+        let reject: (error: Error) => void;
+        let aborted = false;
+        const result = new Promise<R>((f, r) => { fulfill = f; reject = r; });
+
+        const next = () => {
+          if (aborted)
+            return;
+          try {
+            const success = predicate();
+            if (success) {
+              fulfill(success);
               return;
-            try {
-              const success = predicate();
-              if (success) {
-                fulfill(success);
-                return;
-              }
-              if (typeof polling !== 'number')
-                injected.utils.builtins.requestAnimationFrame(next);
-              else
-                injected.utils.builtins.setTimeout(next, polling);
-            } catch (e) {
-              reject(e);
             }
-          };
-
-          next();
-          return { result, abort: () => aborted = true };
-        }, { expression, isFunction, polling: options.pollingInterval, arg });
-        progress.cleanupWhenAborted(() => handle.evaluate(h => h.abort()).catch(() => {}));
-        return handle.evaluateHandle(h => h.result);
-      });
-    }, options.timeout);
+            if (typeof polling !== 'number')
+              injected.utils.builtins.requestAnimationFrame(next);
+            else
+              injected.utils.builtins.setTimeout(next, polling);
+          } catch (e) {
+            reject(e);
+          }
+        };
+
+        next();
+        return { result, abort: () => aborted = true };
+      }, { expression, isFunction, polling: options.pollingInterval, arg }));
+      progress.cleanupWhenAborted(() => handle.evaluate(h => h.abort()).finally(() => handle.dispose()));
+      const result = await progress.race(handle.evaluateHandle(h => h.result));
+      handle.dispose();
+      return result;
+    });
   }
 
   async waitForFunctionValueInUtility<R>(progress: Progress, pageFunction: js.Func1<any, R>) {
@@ -1548,7 +1552,7 @@ export class Frame extends SdkObject {
         return result;
       return JSON.stringify(result);
     }`;
-    const handle = await this._waitForFunctionExpression(serverSideCallMetadata(), expression, true, undefined, { timeout: progress.timeUntilDeadline() }, 'utility');
+    const handle = await this.waitForFunctionExpressionImpl(progress, expression, true, undefined, {}, 'utility');
     return JSON.parse(handle.rawValue()) as R;
   }
 
@@ -1557,18 +1561,18 @@ export class Frame extends SdkObject {
     return context.evaluate(() => document.title);
   }
 
-  async rafrafTimeout(timeout: number): Promise<void> {
+  async rafrafTimeout(progress: Progress, timeout: number): Promise<void> {
     if (timeout === 0)
       return;
-    const context = await this._utilityContext();
+    const context = await progress.race(this._utilityContext());
     await Promise.all([
       // wait for double raf
-      context.evaluate(() => new Promise(x => {
+      progress.race(context.evaluate(() => new Promise(x => {
         requestAnimationFrame(() => {
           requestAnimationFrame(x);
         });
-      })),
-      new Promise(fulfill => setTimeout(fulfill, timeout)),
+      }))),
+      progress.wait(timeout),
     ]);
   }
 

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

@@ -591,12 +591,12 @@ export class Page extends SdkObject {
       return await locator.frame.rafrafTimeoutScreenshotElementWithProgress(progress, locator.selector, timeout, options || {});
     } : async (progress: Progress, timeout: number) => {
       await this.performActionPreChecks(progress);
-      await this.mainFrame().rafrafTimeout(timeout);
+      await this.mainFrame().rafrafTimeout(progress, timeout);
       return await this.screenshotter.screenshotPage(progress, options || {});
     };
 
     const comparator = getComparator('image/png');
-    const controller = new ProgressController(metadata, this);
+    const controller = new ProgressController(metadata, this, 'strict');
     if (!options.expected && options.isNot)
       return { errorMessage: '"not" matcher requires expected result' };
     try {
@@ -632,14 +632,15 @@ export class Page extends SdkObject {
         progress.log(`  generating new stable screenshot expectation`);
       let isFirstIteration = true;
       while (true) {
-        progress.throwIfAborted();
         if (this.isClosed())
           throw new Error('The page has closed');
         const screenshotTimeout = pollIntervals.shift() ?? 1000;
         if (screenshotTimeout)
           progress.log(`waiting ${screenshotTimeout}ms before taking screenshot`);
         previous = actual;
         actual = await rafrafScreenshot(progress, screenshotTimeout).catch(e => {
+          if (isAbortError(e))
+            throw e;
           progress.log(`failed to take screenshot - ` + e.message);
           return undefined;
         });

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

@@ -206,7 +206,6 @@ export class Screenshotter {
       progress.log('taking page screenshot');
       const viewportSize = await this._originalViewportSize(progress);
       await this._preparePageForScreenshot(progress, this._page.mainFrame(), options.style, options.caret !== 'initial', options.animations === 'disabled');
-      progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 
       if (options.fullPage) {
         const fullPageSize = await this._fullPageSize(progress);
@@ -215,14 +214,12 @@ export class Screenshotter {
         if (options.clip)
           documentRect = trimClipToSize(options.clip, documentRect);
         const buffer = await this._screenshot(progress, format, documentRect, undefined, fitsViewport, options);
-        progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
         await this._restorePageAfterScreenshot();
         return buffer;
       }
 
       const viewportRect = options.clip ? trimClipToSize(options.clip, viewportSize) : { x: 0, y: 0, ...viewportSize };
       const buffer = await this._screenshot(progress, format, undefined, viewportRect, true, options);
-      progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
       await this._restorePageAfterScreenshot();
       return buffer;
     });
@@ -235,24 +232,19 @@ export class Screenshotter {
       const viewportSize = await this._originalViewportSize(progress);
 
       await this._preparePageForScreenshot(progress, handle._frame, options.style, options.caret !== 'initial', options.animations === 'disabled');
-      progress.throwIfAborted(); // Do not do extra work.
-
       await handle._waitAndScrollIntoViewIfNeeded(progress, true /* waitForVisible */);
 
-      progress.throwIfAborted(); // Do not do extra work.
-      const boundingBox = await handle.boundingBox();
+      const boundingBox = await progress.race(handle.boundingBox());
       assert(boundingBox, 'Node is either not visible or not an HTMLElement');
       assert(boundingBox.width !== 0, 'Node has 0 width.');
       assert(boundingBox.height !== 0, 'Node has 0 height.');
 
       const fitsViewport = boundingBox.width <= viewportSize.width && boundingBox.height <= viewportSize.height;
-      progress.throwIfAborted(); // Avoid extra work.
       const scrollOffset = await this._page.mainFrame().waitForFunctionValueInUtility(progress, () => ({ x: window.scrollX, y: window.scrollY }));
       const documentRect = { ...boundingBox };
       documentRect.x += scrollOffset.x;
       documentRect.y += scrollOffset.y;
       const buffer = await this._screenshot(progress, format, helper.enclosingIntRect(documentRect), undefined, fitsViewport, options);
-      progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
       await this._restorePageAfterScreenshot();
       return buffer;
     });
@@ -262,71 +254,62 @@ export class Screenshotter {
     if (disableAnimations)
       progress.log('  disabled all CSS animations');
     const syncAnimations = this._page.delegate.shouldToggleStyleSheetToSyncAnimations();
-    await this._page.safeNonStallingEvaluateInAllFrames('(' + inPagePrepareForScreenshots.toString() + `)(${JSON.stringify(screenshotStyle)}, ${hideCaret}, ${disableAnimations}, ${syncAnimations})`, 'utility');
+    progress.cleanupWhenAborted(() => this._restorePageAfterScreenshot());
+    await progress.race(this._page.safeNonStallingEvaluateInAllFrames('(' + inPagePrepareForScreenshots.toString() + `)(${JSON.stringify(screenshotStyle)}, ${hideCaret}, ${disableAnimations}, ${syncAnimations})`, 'utility'));
     if (!process.env.PW_TEST_SCREENSHOT_NO_FONTS_READY) {
       progress.log('waiting for fonts to load...');
-      await frame.nonStallingEvaluateInExistingContext('document.fonts.ready', 'utility').catch(() => {});
+      await progress.race(frame.nonStallingEvaluateInExistingContext('document.fonts.ready', 'utility').catch(() => {}));
       progress.log('fonts loaded');
     }
-    progress.cleanupWhenAborted(() => this._restorePageAfterScreenshot());
   }
 
   async _restorePageAfterScreenshot() {
     await this._page.safeNonStallingEvaluateInAllFrames('window.__pwCleanupScreenshot && window.__pwCleanupScreenshot()', 'utility');
   }
 
   async _maskElements(progress: Progress, options: ScreenshotOptions): Promise<() => Promise<void>> {
+    if (!options.mask || !options.mask.length)
+      return () => Promise.resolve();
+
     const framesToParsedSelectors: MultiMap<Frame, ParsedSelector> = new MultiMap();
 
     const cleanup = async () => {
       await Promise.all([...framesToParsedSelectors.keys()].map(async frame => {
         await frame.hideHighlight();
       }));
     };
+    progress.cleanupWhenAborted(cleanup);
 
-    if (!options.mask || !options.mask.length)
-      return cleanup;
-
-    await Promise.all((options.mask || []).map(async ({ frame, selector }) => {
+    await progress.race(Promise.all((options.mask || []).map(async ({ frame, selector }) => {
       const pair = await frame.selectors.resolveFrameForSelector(selector);
       if (pair)
         framesToParsedSelectors.set(pair.frame, pair.info.parsed);
-    }));
-    progress.throwIfAborted(); // Avoid extra work.
+    })));
 
-    await Promise.all([...framesToParsedSelectors.keys()].map(async frame => {
+    await progress.race(Promise.all([...framesToParsedSelectors.keys()].map(async frame => {
       await frame.maskSelectors(framesToParsedSelectors.get(frame), options.maskColor || '#F0F');
-    }));
-    progress.cleanupWhenAborted(cleanup);
+    })));
     return cleanup;
   }
 
   private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, fitsViewport: boolean, options: ScreenshotOptions): Promise<Buffer> {
     if ((options as any).__testHookBeforeScreenshot)
-      await (options as any).__testHookBeforeScreenshot();
-    progress.throwIfAborted(); // Screenshotting is expensive - avoid extra work.
+      await progress.race((options as any).__testHookBeforeScreenshot());
     const shouldSetDefaultBackground = options.omitBackground && format === 'png';
     if (shouldSetDefaultBackground) {
-      await this._page.delegate.setBackgroundColor({ r: 0, g: 0, b: 0, a: 0 });
       progress.cleanupWhenAborted(() => this._page.delegate.setBackgroundColor());
+      await progress.race(this._page.delegate.setBackgroundColor({ r: 0, g: 0, b: 0, a: 0 }));
     }
-    progress.throwIfAborted(); // Avoid extra work.
 
     const cleanupHighlight = await this._maskElements(progress, options);
-    progress.throwIfAborted(); // Avoid extra work.
-
     const quality = format === 'jpeg' ? options.quality ?? 80 : undefined;
-    const buffer = await this._page.delegate.takeScreenshot(progress, format, documentRect, viewportRect, quality, fitsViewport, options.scale || 'device');
-    progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
-
+    const buffer = await progress.race(this._page.delegate.takeScreenshot(progress, format, documentRect, viewportRect, quality, fitsViewport, options.scale || 'device'));
     await cleanupHighlight();
-    progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
 
     if (shouldSetDefaultBackground)
-      await this._page.delegate.setBackgroundColor();
-    progress.throwIfAborted(); // Avoid side effects.
+      await progress.race(this._page.delegate.setBackgroundColor());
     if ((options as any).__testHookAfterScreenshot)
-      await (options as any).__testHookAfterScreenshot();
+      await progress.race((options as any).__testHookAfterScreenshot());
     return buffer;
   }
 }