chore: add hared TaskQueue for page.screenshot() again (#6714)

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<string, BrowserContext>;
+  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),

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<string | Buffer | void> {
+  async screenshot(options = {}): Promise<string | Buffer> {
     let needsViewportReset = false;
 
     let boundingBox = await this.boundingBox();

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<Buffer | string | void>;
-
-  constructor() {
-    this._chain = Promise.resolve<Buffer | string | void>(undefined);
-  }
-
-  public postTask(
-    task: () => Promise<Buffer | string>
-  ): Promise<Buffer | string | void> {
-    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<Page> {
-    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<string, WebWorker>();
   // 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<Buffer | string | void> {
+  async screenshot(options: ScreenshotOptions = {}): Promise<Buffer | string> {
     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

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<Page>;
   private _workerPromise?: Promise<WebWorker>;
+  private _screenshotTaskQueue: TaskQueue;
   /**
    * @internal
    */
@@ -66,14 +68,16 @@ export class Target {
     browserContext: BrowserContext,
     sessionFactory: () => Promise<CDPSession>,
     ignoreHTTPSErrors: boolean,
-    defaultViewport: Viewport | null
+    defaultViewport: Viewport | null,
+    screenshotTaskQueue: TaskQueue
   ) {
     this._targetInfo = targetInfo;
     this._browserContext = browserContext;
     this._targetId = targetInfo.targetId;
     this._sessionFactory = sessionFactory;
     this._ignoreHTTPSErrors = ignoreHTTPSErrors;
     this._defaultViewport = defaultViewport;
+    this._screenshotTaskQueue = screenshotTaskQueue;
     /** @type {?Promise<!Puppeteer.Page>} */
     this._pagePromise = null;
     /** @type {?Promise<!WebWorker>} */
@@ -121,7 +125,8 @@ export class Target {
           client,
           this,
           this._ignoreHTTPSErrors,
-          this._defaultViewport
+          this._defaultViewport,
+          this._screenshotTaskQueue
         )
       );
     }

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<void>;
+
+  constructor() {
+    this._chain = Promise.resolve();
+  }
+
+  postTask<T>(task: () => Promise<T>): Promise<T> {
+    const result = this._chain.then(task);
+    this._chain = result.then(
+      () => undefined,
+      () => undefined
+    );
+    return result;
+  }
+}

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