From f2e19276acb80e596ff5c781c9ae2bc2f2a8f363 Mon Sep 17 00:00:00 2001 From: dmitrysteblyuk Date: Thu, 23 Sep 2021 14:37:35 +0200 Subject: [PATCH] chore: add hared TaskQueue for `page.screenshot()` again (#6714) --- src/common/Browser.ts | 6 +++++- src/common/JSHandle.ts | 2 +- src/common/Page.ts | 42 ++++++++++++++++++----------------------- src/common/Target.ts | 9 +++++++-- src/common/TaskQueue.ts | 32 +++++++++++++++++++++++++++++++ test/headful.spec.ts | 32 +++++++++++++++++++++++++++++++ 6 files changed, 95 insertions(+), 28 deletions(-) create mode 100644 src/common/TaskQueue.ts diff --git a/src/common/Browser.ts b/src/common/Browser.ts index 49669f1e713..e8ff50bb601 100644 --- a/src/common/Browser.ts +++ b/src/common/Browser.ts @@ -21,6 +21,7 @@ import { EventEmitter } from './EventEmitter.js'; import { Connection, ConnectionEmittedEvents } from './Connection.js'; import { Protocol } from 'devtools-protocol'; import { Page } from './Page.js'; +import { TaskQueue } from './TaskQueue.js'; import { ChildProcess } from 'child_process'; import { Viewport } from './PuppeteerViewport.js'; @@ -238,6 +239,7 @@ export class Browser extends EventEmitter { private _targetFilterCallback: TargetFilterCallback; private _defaultContext: BrowserContext; private _contexts: Map; + private _screenshotTaskQueue: TaskQueue; /** * @internal * Used in Target.ts directly so cannot be marked private. @@ -260,6 +262,7 @@ export class Browser extends EventEmitter { this._ignoreHTTPSErrors = ignoreHTTPSErrors; this._defaultViewport = defaultViewport; this._process = process; + this._screenshotTaskQueue = new TaskQueue(); this._connection = connection; this._closeCallback = closeCallback || function (): void {}; this._targetFilterCallback = targetFilterCallback || ((): boolean => true); @@ -379,7 +382,8 @@ export class Browser extends EventEmitter { context, () => this._connection.createSession(targetInfo), this._ignoreHTTPSErrors, - this._defaultViewport + this._defaultViewport, + this._screenshotTaskQueue ); assert( !this._targets.has(event.targetInfo.targetId), diff --git a/src/common/JSHandle.ts b/src/common/JSHandle.ts index 347f1943580..6dbb2537098 100644 --- a/src/common/JSHandle.ts +++ b/src/common/JSHandle.ts @@ -813,7 +813,7 @@ export class ElementHandle< * {@link Page.screenshot} to take a screenshot of the element. * If the element is detached from DOM, the method throws an error. */ - async screenshot(options = {}): Promise { + async screenshot(options = {}): Promise { let needsViewportReset = false; let boundingBox = await this.boundingBox(); diff --git a/src/common/Page.ts b/src/common/Page.ts index 9ebf8a3376e..a315840bdec 100644 --- a/src/common/Page.ts +++ b/src/common/Page.ts @@ -62,6 +62,7 @@ import { } from './EvalTypes.js'; import { PDFOptions, paperFormats } from './PDFOptions.js'; import { isNode } from '../environment.js'; +import { TaskQueue } from './TaskQueue.js'; /** * @public @@ -370,22 +371,6 @@ export interface PageEventObject { workerdestroyed: WebWorker; } -class ScreenshotTaskQueue { - _chain: Promise; - - constructor() { - this._chain = Promise.resolve(undefined); - } - - public postTask( - task: () => Promise - ): Promise { - const result = this._chain.then(task); - this._chain = result.catch(() => {}); - return result; - } -} - /** * Page provides methods to interact with a single tab or * {@link https://developer.chrome.com/extensions/background_pages | extension background page} in Chromium. @@ -437,9 +422,15 @@ export class Page extends EventEmitter { client: CDPSession, target: Target, ignoreHTTPSErrors: boolean, - defaultViewport: Viewport | null + defaultViewport: Viewport | null, + screenshotTaskQueue: TaskQueue ): Promise { - const page = new Page(client, target, ignoreHTTPSErrors); + const page = new Page( + client, + target, + ignoreHTTPSErrors, + screenshotTaskQueue + ); await page._initialize(); if (defaultViewport) await page.setViewport(defaultViewport); return page; @@ -460,7 +451,7 @@ export class Page extends EventEmitter { private _coverage: Coverage; private _javascriptEnabled = true; private _viewport: Viewport | null; - private _screenshotTaskQueue: ScreenshotTaskQueue; + private _screenshotTaskQueue: TaskQueue; private _workers = new Map(); // TODO: improve this typedef - it's a function that takes a file chooser or // something? @@ -472,7 +463,12 @@ export class Page extends EventEmitter { /** * @internal */ - constructor(client: CDPSession, target: Target, ignoreHTTPSErrors: boolean) { + constructor( + client: CDPSession, + target: Target, + ignoreHTTPSErrors: boolean, + screenshotTaskQueue: TaskQueue + ) { super(); this._client = client; this._target = target; @@ -489,7 +485,7 @@ export class Page extends EventEmitter { this._emulationManager = new EmulationManager(client); this._tracing = new Tracing(client); this._coverage = new Coverage(client); - this._screenshotTaskQueue = new ScreenshotTaskQueue(); + this._screenshotTaskQueue = screenshotTaskQueue; this._viewport = null; client.on('Target.attachedToTarget', (event) => { @@ -2552,9 +2548,7 @@ export class Page extends EventEmitter { * @returns Promise which resolves to buffer or a base64 string (depending on * the value of `encoding`) with captured screenshot. */ - async screenshot( - options: ScreenshotOptions = {} - ): Promise { + async screenshot(options: ScreenshotOptions = {}): Promise { let screenshotType = null; // options.type takes precedence over inferring the type from options.path // because it may be a 0-length file with no extension created beforehand diff --git a/src/common/Target.ts b/src/common/Target.ts index 42506f1b1fa..8e890a603c9 100644 --- a/src/common/Target.ts +++ b/src/common/Target.ts @@ -20,6 +20,7 @@ import { CDPSession } from './Connection.js'; import { Browser, BrowserContext } from './Browser.js'; import { Viewport } from './PuppeteerViewport.js'; import { Protocol } from 'devtools-protocol'; +import { TaskQueue } from './TaskQueue.js'; /** * @public @@ -33,6 +34,7 @@ export class Target { private _defaultViewport?: Viewport; private _pagePromise?: Promise; private _workerPromise?: Promise; + private _screenshotTaskQueue: TaskQueue; /** * @internal */ @@ -66,7 +68,8 @@ export class Target { browserContext: BrowserContext, sessionFactory: () => Promise, ignoreHTTPSErrors: boolean, - defaultViewport: Viewport | null + defaultViewport: Viewport | null, + screenshotTaskQueue: TaskQueue ) { this._targetInfo = targetInfo; this._browserContext = browserContext; @@ -74,6 +77,7 @@ export class Target { this._sessionFactory = sessionFactory; this._ignoreHTTPSErrors = ignoreHTTPSErrors; this._defaultViewport = defaultViewport; + this._screenshotTaskQueue = screenshotTaskQueue; /** @type {?Promise} */ this._pagePromise = null; /** @type {?Promise} */ @@ -121,7 +125,8 @@ export class Target { client, this, this._ignoreHTTPSErrors, - this._defaultViewport + this._defaultViewport, + this._screenshotTaskQueue ) ); } diff --git a/src/common/TaskQueue.ts b/src/common/TaskQueue.ts new file mode 100644 index 00000000000..bf1ffbc00c2 --- /dev/null +++ b/src/common/TaskQueue.ts @@ -0,0 +1,32 @@ +/** + * Copyright 2020 Google Inc. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +export class TaskQueue { + private _chain: Promise; + + constructor() { + this._chain = Promise.resolve(); + } + + postTask(task: () => Promise): Promise { + const result = this._chain.then(task); + this._chain = result.then( + () => undefined, + () => undefined + ); + return result; + } +} diff --git a/test/headful.spec.ts b/test/headful.spec.ts index 62441450c43..57bacacef11 100644 --- a/test/headful.spec.ts +++ b/test/headful.spec.ts @@ -292,4 +292,36 @@ describeChromeOnly('headful tests', function () { await browser.close(); }); }); + + describe('Page.screenshot', function () { + it('should run in parallel in multiple pages', async () => { + const { server, puppeteer } = getTestState(); + const browser = await puppeteer.launch(headfulOptions); + const context = await browser.createIncognitoBrowserContext(); + + const N = 2; + const pages = await Promise.all( + Array(N) + .fill(0) + .map(async () => { + const page = await context.newPage(); + await page.goto(server.PREFIX + '/grid.html'); + return page; + }) + ); + const promises = []; + for (let i = 0; i < N; ++i) + promises.push( + pages[i].screenshot({ + clip: { x: 50 * i, y: 0, width: 50, height: 50 }, + }) + ); + const screenshots = await Promise.all(promises); + for (let i = 0; i < N; ++i) + expect(screenshots[i]).toBeGolden(`grid-cell-${i}.png`); + await Promise.all(pages.map((page) => page.close())); + + await browser.close(); + }); + }); });