fix(screenshot): never throw page is navigating (#6103)

src/server/chromium/crExecutionContext.ts

@@ -69,10 +69,7 @@ export class CRExecutionContext implements js.ExecutionContextDelegate {
     return returnByValue ? parseEvaluationResultValue(remoteObject.value) : utilityScript._context.createHandle(remoteObject);
   }
 
-  async getProperties(handle: js.JSHandle): Promise<Map<string, js.JSHandle>> {
-    const objectId = handle._objectId;
-    if (!objectId)
-      return new Map();
+  async getProperties(context: js.ExecutionContext, objectId: js.ObjectId): Promise<Map<string, js.JSHandle>> {
     const response = await this._client.send('Runtime.getProperties', {
       objectId,
       ownProperties: true
@@ -81,7 +78,7 @@ export class CRExecutionContext implements js.ExecutionContextDelegate {
     for (const property of response.result) {
       if (!property.enumerable || !property.value)
         continue;
-      result.set(property.name, handle._context.createHandle(property.value));
+      result.set(property.name, context.createHandle(property.value));
     }
     return result;
   }
@@ -90,10 +87,8 @@ export class CRExecutionContext implements js.ExecutionContextDelegate {
     return new js.JSHandle(context, remoteObject.subtype || remoteObject.type, remoteObject.objectId, potentiallyUnserializableValue(remoteObject));
   }
 
-  async releaseHandle(handle: js.JSHandle): Promise<void> {
-    if (!handle._objectId)
-      return;
-    await releaseObject(this._client, handle._objectId);
+  async releaseHandle(objectId: js.ObjectId): Promise<void> {
+    await releaseObject(this._client, objectId);
   }
 }
 

src/server/chromium/crPage.ts

@@ -38,6 +38,7 @@ import { ConsoleMessage } from '../console';
 import { rewriteErrorMessage } from '../../utils/stackTrace';
 import { assert, headersArrayToObject, createGuid, canAccessFile } from '../../utils/utils';
 import { VideoRecorder } from './videoRecorder';
+import { Progress } from '../progress';
 
 
 const UTILITY_WORLD_NAME = '__playwright_utility_world__';
@@ -245,7 +246,7 @@ export class CRPage implements PageDelegate {
     await this._mainFrameSession._client.send('Emulation.setDefaultBackgroundColorOverride', { color });
   }
 
-  async takeScreenshot(format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer> {
+  async takeScreenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer> {
     const { visualViewport } = await this._mainFrameSession._client.send('Page.getLayoutMetrics');
     if (!documentRect) {
       documentRect = {
@@ -260,6 +261,7 @@ export class CRPage implements PageDelegate {
     // When taking screenshots with documentRect (based on the page content, not viewport),
     // ignore current page scale.
     const clip = { ...documentRect, scale: viewportRect ? visualViewport.scale : 1 };
+    progress.throwIfAborted();
     const result = await this._mainFrameSession._client.send('Page.captureScreenshot', { format, quality, clip });
     return Buffer.from(result.data, 'base64');
   }

src/server/dom.ts

@@ -51,7 +51,7 @@ export class FrameExecutionContext extends js.ExecutionContext {
     return js.evaluate(this, true /* returnByValue */, pageFunction, arg);
   }
 
-  async evaluateHandle<Arg, R>(pageFunction: js.Func1<Arg, R>, arg: Arg): Promise<js.SmartHandle<R>> {
+  async evaluateHandle<Arg, R>(pageFunction: js.Func1<Arg, R>, arg?: Arg): Promise<js.SmartHandle<R>> {
     return js.evaluate(this, false /* returnByValue */, pageFunction, arg);
   }
 

src/server/firefox/ffExecutionContext.ts

@@ -66,17 +66,14 @@ export class FFExecutionContext implements js.ExecutionContextDelegate {
     return utilityScript._context.createHandle(payload.result!);
   }
 
-  async getProperties(handle: js.JSHandle): Promise<Map<string, js.JSHandle>> {
-    const objectId = handle._objectId;
-    if (!objectId)
-      return new Map();
+  async getProperties(context: js.ExecutionContext, objectId: js.ObjectId): Promise<Map<string, js.JSHandle>> {
     const response = await this._session.send('Runtime.getObjectProperties', {
       executionContextId: this._executionContextId,
       objectId,
     });
     const result = new Map();
     for (const property of response.properties)
-      result.set(property.name, handle._context.createHandle(property.value));
+      result.set(property.name, context.createHandle(property.value));
     return result;
   }
 
@@ -84,12 +81,10 @@ export class FFExecutionContext implements js.ExecutionContextDelegate {
     return new js.JSHandle(context, remoteObject.subtype || remoteObject.type || '', remoteObject.objectId, potentiallyUnserializableValue(remoteObject));
   }
 
-  async releaseHandle(handle: js.JSHandle): Promise<void> {
-    if (!handle._objectId)
-      return;
+  async releaseHandle(objectId: js.ObjectId): Promise<void> {
     await this._session.send('Runtime.disposeObject', {
       executionContextId: this._executionContextId,
-      objectId: handle._objectId,
+      objectId
     });
   }
 }

src/server/firefox/ffPage.ts

@@ -21,7 +21,6 @@ import * as frames from '../frames';
 import { helper, RegisteredListener } from '../helper';
 import { assert } from '../../utils/utils';
 import { Page, PageBinding, PageDelegate, Worker } from '../page';
-import { kScreenshotDuringNavigationError } from '../screenshotter';
 import * as types from '../types';
 import { getAccessibilityTree } from './ffAccessibility';
 import { FFBrowserContext } from './ffBrowser';
@@ -30,7 +29,7 @@ import { FFExecutionContext } from './ffExecutionContext';
 import { RawKeyboardImpl, RawMouseImpl, RawTouchscreenImpl } from './ffInput';
 import { FFNetworkManager } from './ffNetworkManager';
 import { Protocol } from './protocol';
-import { rewriteErrorMessage } from '../../utils/stackTrace';
+import { Progress } from '../progress';
 
 const UTILITY_WORLD_NAME = '__playwright_utility_world__';
 
@@ -394,10 +393,9 @@ export class FFPage implements PageDelegate {
       throw new Error('Not implemented');
   }
 
-  async takeScreenshot(format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer> {
+  async takeScreenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer> {
     if (!documentRect) {
-      const context = await this._page.mainFrame()._utilityContext();
-      const scrollOffset = await context.evaluate(() => ({ x: window.scrollX, y: window.scrollY }));
+      const scrollOffset = await this._page.mainFrame().waitForFunctionValue(progress, () => ({ x: window.scrollX, y: window.scrollY }));
       documentRect = {
         x: viewportRect!.x + scrollOffset.x,
         y: viewportRect!.y + scrollOffset.y,
@@ -407,13 +405,10 @@ export class FFPage implements PageDelegate {
     }
     // TODO: remove fullPage option from Page.screenshot.
     // TODO: remove Page.getBoundingBox method.
+    progress.throwIfAborted();
     const { data } = await this._session.send('Page.screenshot', {
       mimeType: ('image/' + format) as ('image/png' | 'image/jpeg'),
       clip: documentRect,
-    }).catch(e => {
-      if (e instanceof Error && e.message.includes('document.documentElement is null'))
-        rewriteErrorMessage(e, kScreenshotDuringNavigationError);
-      throw e;
     });
     return Buffer.from(data, 'base64');
   }

src/server/frames.ts

@@ -26,7 +26,7 @@ import { BrowserContext } from './browserContext';
 import { Progress, ProgressController } from './progress';
 import { assert, createGuid, makeWaitForNextTask } from '../utils/utils';
 import { debugLogger } from '../utils/debugLogger';
-import { CallMetadata, SdkObject } from './instrumentation';
+import { CallMetadata, internalCallMetadata, SdkObject } from './instrumentation';
 import { ElementStateWithoutStable } from './injected/injectedScript';
 
 type ContextData = {
@@ -1081,6 +1081,18 @@ export class Frame extends SdkObject {
         this._page._timeoutSettings.timeout(options));
   }
 
+  async waitForFunctionValue<R>(progress: Progress, pageFunction: js.Func1<any, R>) {
+    const expression = `() => {
+      const result = (${pageFunction})();
+      if (!result)
+        return result;
+      return JSON.stringify(result);
+
+    }`;
+    const handle = await this._page.mainFrame()._waitForFunctionExpression(internalCallMetadata(), expression, true, undefined, { timeout: progress.timeUntilDeadline() });
+    return JSON.parse(handle.rawValue()) as R;
+  }
+
   async title(): Promise<string> {
     const context = await this._utilityContext();
     return context.evaluate(() => document.title);

src/server/javascript.ts

@@ -20,7 +20,7 @@ import { serializeAsCallArgument } from './common/utilityScriptSerializers';
 import type UtilityScript from './injected/utilityScript';
 import { SdkObject } from './instrumentation';
 
-type ObjectId = string;
+export type ObjectId = string;
 export type RemoteObject = {
   objectId?: ObjectId,
   value?: any
@@ -46,9 +46,9 @@ export interface ExecutionContextDelegate {
   rawEvaluate(expression: string): Promise<ObjectId>;
   rawCallFunctionNoReply(func: Function, ...args: any[]): void;
   evaluateWithArguments(expression: string, returnByValue: boolean, utilityScript: JSHandle<any>, values: any[], objectIds: ObjectId[]): Promise<any>;
-  getProperties(handle: JSHandle): Promise<Map<string, JSHandle>>;
+  getProperties(context: ExecutionContext, objectId: ObjectId): Promise<Map<string, JSHandle>>;
   createHandle(context: ExecutionContext, remoteObject: RemoteObject): JSHandle;
-  releaseHandle(handle: JSHandle): Promise<void>;
+  releaseHandle(objectId: ObjectId): Promise<void>;
 }
 
 export class ExecutionContext extends SdkObject {
@@ -144,8 +144,14 @@ export class JSHandle<T = any> extends SdkObject {
     return result;
   }
 
-  getProperties(): Promise<Map<string, JSHandle>> {
-    return this._context._delegate.getProperties(this);
+  async getProperties(): Promise<Map<string, JSHandle>> {
+    if (!this._objectId)
+      return new Map();
+    return this._context._delegate.getProperties(this._context, this._objectId);
+  }
+
+  rawValue() {
+    return this._value;
   }
 
   async jsonValue(): Promise<T> {
@@ -164,7 +170,8 @@ export class JSHandle<T = any> extends SdkObject {
     if (this._disposed)
       return;
     this._disposed = true;
-    this._context._delegate.releaseHandle(this).catch(e => {});
+    if (this._objectId)
+      this._context._delegate.releaseHandle(this._objectId).catch(e => {});
   }
 
   toString(): string {

src/server/page.ts

@@ -27,7 +27,7 @@ import { BrowserContext } from './browserContext';
 import { ConsoleMessage } from './console';
 import * as accessibility from './accessibility';
 import { FileChooser } from './fileChooser';
-import { ProgressController } from './progress';
+import { Progress, ProgressController } from './progress';
 import { assert, createGuid, isError } from '../utils/utils';
 import { debugLogger } from '../utils/debugLogger';
 import { Selectors } from './selectors';
@@ -59,7 +59,7 @@ export interface PageDelegate {
   canScreenshotOutsideViewport(): boolean;
   resetViewport(): Promise<void>; // Only called if canScreenshotOutsideViewport() returns false.
   setBackgroundColor(color?: { r: number; g: number; b: number; a: number; }): Promise<void>;
-  takeScreenshot(format: string, documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer>;
+  takeScreenshot(progress: Progress, format: string, documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer>;
 
   isElementHandle(remoteObject: any): boolean;
   adoptElementHandle<T extends Node>(handle: dom.ElementHandle<T>, to: dom.FrameExecutionContext): Promise<dom.ElementHandle<T>>;

src/server/screenshotter.ts

@@ -19,7 +19,6 @@ import * as dom from './dom';
 import { helper } from './helper';
 import { Page } from './page';
 import * as types from './types';
-import { rewriteErrorMessage } from '../utils/stackTrace';
 import { Progress } from './progress';
 import { assert } from '../utils/utils';
 
@@ -32,19 +31,16 @@ export class Screenshotter {
     this._queue = new TaskQueue();
   }
 
-  private async _originalViewportSize(): Promise<{ viewportSize: types.Size, originalViewportSize: types.Size | null }> {
+  private async _originalViewportSize(progress: Progress): Promise<{ viewportSize: types.Size, originalViewportSize: types.Size | null }> {
     const originalViewportSize = this._page.viewportSize();
     let viewportSize = originalViewportSize;
-    if (!viewportSize) {
-      const context = await this._page.mainFrame()._utilityContext();
-      viewportSize = await context.evaluate(() => ({ width: window.innerWidth, height: window.innerHeight }));
-    }
+    if (!viewportSize)
+      viewportSize = await this._page.mainFrame().waitForFunctionValue(progress, () => ({ width: window.innerWidth, height: window.innerHeight }));
     return { viewportSize, originalViewportSize };
   }
 
-  private async _fullPageSize(): Promise<types.Size> {
-    const context = await this._page.mainFrame()._utilityContext();
-    const fullPageSize = await context.evaluate(() => {
+  private async _fullPageSize(progress: Progress): Promise<types.Size> {
+    const fullPageSize = await this._page.mainFrame().waitForFunctionValue(progress, () => {
       if (!document.body || !document.documentElement)
         return null;
       return {
@@ -60,18 +56,16 @@ export class Screenshotter {
         ),
       };
     });
-    if (!fullPageSize)
-      throw new Error(kScreenshotDuringNavigationError);
-    return fullPageSize;
+    return fullPageSize!;
   }
 
   async screenshotPage(progress: Progress, options: types.ScreenshotOptions): Promise<Buffer> {
     const format = validateScreenshotOptions(options);
     return this._queue.postTask(async () => {
-      const { viewportSize, originalViewportSize } = await this._originalViewportSize();
+      const { viewportSize, originalViewportSize } = await this._originalViewportSize(progress);
 
       if (options.fullPage) {
-        const fullPageSize = await this._fullPageSize();
+        const fullPageSize = await this._fullPageSize(progress);
         let documentRect = { x: 0, y: 0, width: fullPageSize.width, height: fullPageSize.height };
         let overriddenViewportSize: types.Size | null = null;
         const fitsViewport = fullPageSize.width <= viewportSize.width && fullPageSize.height <= viewportSize.height;
@@ -92,15 +86,17 @@ export class Screenshotter {
 
       const viewportRect = options.clip ? trimClipToSize(options.clip, viewportSize) : { x: 0, y: 0, ...viewportSize };
       return await this._screenshot(progress, format, undefined, viewportRect, options);
-    }).catch(rewriteError);
+    });
   }
 
   async screenshotElement(progress: Progress, handle: dom.ElementHandle, options: types.ElementScreenshotOptions = {}): Promise<Buffer> {
     const format = validateScreenshotOptions(options);
     return this._queue.postTask(async () => {
-      const { viewportSize, originalViewportSize } = await this._originalViewportSize();
+      const { viewportSize, originalViewportSize } = await this._originalViewportSize(progress);
 
       await handle._waitAndScrollIntoViewIfNeeded(progress);
+
+      progress.throwIfAborted(); // Do not do extra work.
       let boundingBox = await handle.boundingBox();
       assert(boundingBox, 'Node is either not visible or not an HTMLElement');
       assert(boundingBox.width !== 0, 'Node has 0 width.');
@@ -117,6 +113,7 @@ export class Screenshotter {
         await this._page.setViewportSize(overriddenViewportSize);
         progress.cleanupWhenAborted(() => this._restoreViewport(originalViewportSize));
 
+        progress.throwIfAborted(); // Avoid extra work.
         await handle._waitAndScrollIntoViewIfNeeded(progress);
         boundingBox = await handle.boundingBox();
         assert(boundingBox, 'Node is either not visible or not an HTMLElement');
@@ -124,8 +121,8 @@ export class Screenshotter {
         assert(boundingBox.height !== 0, 'Node has 0 height.');
       }
 
-      const context = await this._page.mainFrame()._utilityContext();
-      const scrollOffset = await context.evaluate(() => ({ x: window.scrollX, y: window.scrollY }));
+      progress.throwIfAborted(); // Avoid extra work.
+      const scrollOffset = await this._page.mainFrame().waitForFunctionValue(progress, () => ({ x: window.scrollX, y: window.scrollY }));
       const documentRect = { ...boundingBox };
       documentRect.x += scrollOffset.x;
       documentRect.y += scrollOffset.y;
@@ -134,7 +131,7 @@ export class Screenshotter {
       if (overriddenViewportSize)
         await this._restoreViewport(originalViewportSize);
       return buffer;
-    }).catch(rewriteError);
+    });
   }
 
   private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, options: types.ElementScreenshotOptions): Promise<Buffer> {
@@ -146,7 +143,8 @@ export class Screenshotter {
       await this._page._delegate.setBackgroundColor({ r: 0, g: 0, b: 0, a: 0});
       progress.cleanupWhenAborted(() => this._page._delegate.setBackgroundColor());
     }
-    const buffer = await this._page._delegate.takeScreenshot(format, documentRect, viewportRect, options.quality);
+    progress.throwIfAborted(); // Avoid extra work.
+    const buffer = await this._page._delegate.takeScreenshot(progress, format, documentRect, viewportRect, options.quality);
     progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
     if (shouldSetDefaultBackground)
       await this._page._delegate.setBackgroundColor();
@@ -221,10 +219,3 @@ function validateScreenshotOptions(options: types.ScreenshotOptions): 'png' | 'j
   }
   return format;
 }
-
-export const kScreenshotDuringNavigationError = 'Cannot take a screenshot while page is navigating';
-function rewriteError(e: any) {
-  if (typeof e === 'object' && e instanceof Error && e.message.includes('Execution context was destroyed'))
-    rewriteErrorMessage(e, kScreenshotDuringNavigationError);
-  throw e;
-}

src/server/webkit/wkExecutionContext.ts

@@ -118,10 +118,7 @@ export class WKExecutionContext implements js.ExecutionContextDelegate {
     }
   }
 
-  async getProperties(handle: js.JSHandle): Promise<Map<string, js.JSHandle>> {
-    const objectId = handle._objectId;
-    if (!objectId)
-      return new Map();
+  async getProperties(context: js.ExecutionContext, objectId: js.ObjectId): Promise<Map<string, js.JSHandle>> {
     const response = await this._session.send('Runtime.getProperties', {
       objectId,
       ownProperties: true
@@ -130,7 +127,7 @@ export class WKExecutionContext implements js.ExecutionContextDelegate {
     for (const property of response.properties) {
       if (!property.enumerable || !property.value)
         continue;
-      result.set(property.name, handle._context.createHandle(property.value));
+      result.set(property.name, context.createHandle(property.value));
     }
     return result;
   }
@@ -140,10 +137,8 @@ export class WKExecutionContext implements js.ExecutionContextDelegate {
     return new js.JSHandle(context, isPromise ? 'promise' : remoteObject.subtype || remoteObject.type, remoteObject.objectId, potentiallyUnserializableValue(remoteObject));
   }
 
-  async releaseHandle(handle: js.JSHandle): Promise<void> {
-    if (!handle._objectId)
-      return;
-    await this._session.send('Runtime.releaseObject', {objectId: handle._objectId});
+  async releaseHandle(objectId: js.ObjectId): Promise<void> {
+    await this._session.send('Runtime.releaseObject', { objectId });
   }
 }
 

src/server/webkit/wkPage.ts

@@ -27,6 +27,7 @@ import { helper, RegisteredListener } from '../helper';
 import { JSHandle } from '../javascript';
 import * as network from '../network';
 import { Page, PageBinding, PageDelegate } from '../page';
+import { Progress } from '../progress';
 import * as types from '../types';
 import { Protocol } from './protocol';
 import { getAccessibilityTree } from './wkAccessibility';
@@ -750,7 +751,7 @@ export class WKPage implements PageDelegate {
     this._recordingVideoFile = null;
   }
 
-  async takeScreenshot(format: string, documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer> {
+  async takeScreenshot(progress: Progress, format: string, documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, quality: number | undefined): Promise<Buffer> {
     const rect = (documentRect || viewportRect)!;
     const result = await this._session.send('Page.snapshotRect', { ...rect, coordinateSystem: documentRect ? 'Page' : 'Viewport' });
     const prefix = 'data:image/png;base64,';

tests/page-screenshot.spec.ts

@@ -368,4 +368,21 @@ browserTest.describe('page screenshot', () => {
     expect(pixel(0, 8339).r).toBeLessThan(128);
     expect(pixel(0, 8339).b).toBeGreaterThan(128);
   });
+
+  it('should take fullPage screenshots during navigation', async ({page, server}) => {
+    await page.setViewportSize({width: 500, height: 500});
+    await page.goto(server.PREFIX + '/grid.html');
+    const reloadSeveralTimes = async () => {
+      for (let i = 0; i < 5; ++i)
+        await page.reload();
+    };
+    const screenshotSeveralTimes = async () => {
+      for (let i = 0; i < 5; ++i)
+        await page.screenshot({ fullPage: true });
+    };
+    await Promise.all([
+      reloadSeveralTimes(),
+      screenshotSeveralTimes()
+    ]);
+  });
 });