1 files changed, 255 insertions, 0 deletions
diff --git a/remote/test/puppeteer/packages/puppeteer-core/src/node/ScreenRecorder.ts b/remote/test/puppeteer/packages/puppeteer-core/src/node/ScreenRecorder.ts
new file mode 100644
index 0000000000..effb2d63ba
--- /dev/null
+++ b/remote/test/puppeteer/packages/puppeteer-core/src/node/ScreenRecorder.ts
@@ -0,0 +1,255 @@
+/**
+ * @license
+ * Copyright 2023 Google Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type {ChildProcessWithoutNullStreams} from 'child_process';
+import {spawn, spawnSync} from 'child_process';
+import {PassThrough} from 'stream';
+
+import debug from 'debug';
+
+import type {OperatorFunction} from '../../third_party/rxjs/rxjs.js';
+import {
+  bufferCount,
+  concatMap,
+  filter,
+  from,
+  fromEvent,
+  lastValueFrom,
+  map,
+  takeUntil,
+  tap,
+} from '../../third_party/rxjs/rxjs.js';
+import {CDPSessionEvent} from '../api/CDPSession.js';
+import type {BoundingBox} from '../api/ElementHandle.js';
+import type {Page} from '../api/Page.js';
+import {debugError, fromEmitterEvent} from '../common/util.js';
+import {guarded} from '../util/decorators.js';
+import {asyncDisposeSymbol} from '../util/disposable.js';
+
+const CRF_VALUE = 30;
+const DEFAULT_FPS = 30;
+
+const debugFfmpeg = debug('puppeteer:ffmpeg');
+
+/**
+ * @internal
+ */
+export interface ScreenRecorderOptions {
+  speed?: number;
+  crop?: BoundingBox;
+  format?: 'gif' | 'webm';
+  scale?: number;
+  path?: string;
+}
+
+/**
+ * @public
+ */
+export class ScreenRecorder extends PassThrough {
+  #page: Page;
+
+  #process: ChildProcessWithoutNullStreams;
+
+  #controller = new AbortController();
+  #lastFrame: Promise<readonly [Buffer, number]>;
+
+  /**
+   * @internal
+   */
+  constructor(
+    page: Page,
+    width: number,
+    height: number,
+    {speed, scale, crop, format, path}: ScreenRecorderOptions = {}
+  ) {
+    super({allowHalfOpen: false});
+
+    path ??= 'ffmpeg';
+
+    // Tests if `ffmpeg` exists.
+    const {error} = spawnSync(path);
+    if (error) {
+      throw error;
+    }
+
+    this.#process = spawn(
+      path,
+      // See https://trac.ffmpeg.org/wiki/Encode/VP9 for more information on flags.
+      [
+        ['-loglevel', 'error'],
+        // Reduces general buffering.
+        ['-avioflags', 'direct'],
+        // Reduces initial buffering while analyzing input fps and other stats.
+        [
+          '-fpsprobesize',
+          '0',
+          '-probesize',
+          '32',
+          '-analyzeduration',
+          '0',
+          '-fflags',
+          'nobuffer',
+        ],
+        // Forces input to be read from standard input, and forces png input
+        // image format.
+        ['-f', 'image2pipe', '-c:v', 'png', '-i', 'pipe:0'],
+        // Overwrite output and no audio.
+        ['-y', '-an'],
+        // This drastically reduces stalling when cpu is overbooked. By default
+        // VP9 tries to use all available threads?
+        ['-threads', '1'],
+        // Specifies the frame rate we are giving ffmpeg.
+        ['-framerate', `${DEFAULT_FPS}`],
+        // Specifies the encoding and format we are using.
+        this.#getFormatArgs(format ?? 'webm'),
+        // Disable bitrate.
+        ['-b:v', '0'],
+        // Filters to ensure the images are piped correctly.
+        [
+          '-vf',
+          `${
+            speed ? `setpts=${1 / speed}*PTS,` : ''
+          }crop='min(${width},iw):min(${height},ih):0:0',pad=${width}:${height}:0:0${
+            crop ? `,crop=${crop.width}:${crop.height}:${crop.x}:${crop.y}` : ''
+          }${scale ? `,scale=iw*${scale}:-1` : ''}`,
+        ],
+        'pipe:1',
+      ].flat(),
+      {stdio: ['pipe', 'pipe', 'pipe']}
+    );
+    this.#process.stdout.pipe(this);
+    this.#process.stderr.on('data', (data: Buffer) => {
+      debugFfmpeg(data.toString('utf8'));
+    });
+
+    this.#page = page;
+
+    const {client} = this.#page.mainFrame();
+    client.once(CDPSessionEvent.Disconnected, () => {
+      void this.stop().catch(debugError);
+    });
+
+    this.#lastFrame = lastValueFrom(
+      fromEmitterEvent(client, 'Page.screencastFrame').pipe(
+        tap(event => {
+          void client.send('Page.screencastFrameAck', {
+            sessionId: event.sessionId,
+          });
+        }),
+        filter(event => {
+          return event.metadata.timestamp !== undefined;
+        }),
+        map(event => {
+          return {
+            buffer: Buffer.from(event.data, 'base64'),
+            timestamp: event.metadata.timestamp!,
+          };
+        }),
+        bufferCount(2, 1) as OperatorFunction<
+          {buffer: Buffer; timestamp: number},
+          [
+            {buffer: Buffer; timestamp: number},
+            {buffer: Buffer; timestamp: number},
+          ]
+        >,
+        concatMap(([{timestamp: previousTimestamp, buffer}, {timestamp}]) => {
+          return from(
+            Array<Buffer>(
+              Math.round(
+                DEFAULT_FPS * Math.max(timestamp - previousTimestamp, 0)
+              )
+            ).fill(buffer)
+          );
+        }),
+        map(buffer => {
+          void this.#writeFrame(buffer);
+          return [buffer, performance.now()] as const;
+        }),
+        takeUntil(fromEvent(this.#controller.signal, 'abort'))
+      ),
+      {defaultValue: [Buffer.from([]), performance.now()] as const}
+    );
+  }
+
+  #getFormatArgs(format: 'webm' | 'gif') {
+    switch (format) {
+      case 'webm':
+        return [
+          // Sets the codec to use.
+          ['-c:v', 'vp9'],
+          // Sets the format
+          ['-f', 'webm'],
+          // Sets the quality. Lower the better.
+          ['-crf', `${CRF_VALUE}`],
+          // Sets the quality and how efficient the compression will be.
+          ['-deadline', 'realtime', '-cpu-used', '8'],
+        ].flat();
+      case 'gif':
+        return [
+          // Sets the frame rate and uses a custom palette generated from the
+          // input.
+          [
+            '-vf',
+            'fps=5,split[s0][s1];[s0]palettegen=stats_mode=diff[p];[s1][p]paletteuse',
+          ],
+          // Sets the format
+          ['-f', 'gif'],
+        ].flat();
+    }
+  }
+
+  @guarded()
+  async #writeFrame(buffer: Buffer) {
+    const error = await new Promise<Error | null | undefined>(resolve => {
+      this.#process.stdin.write(buffer, resolve);
+    });
+    if (error) {
+      console.log(`ffmpeg failed to write: ${error.message}.`);
+    }
+  }
+
+  /**
+   * Stops the recorder.
+   *
+   * @public
+   */
+  @guarded()
+  async stop(): Promise<void> {
+    if (this.#controller.signal.aborted) {
+      return;
+    }
+    // Stopping the screencast will flush the frames.
+    await this.#page._stopScreencast().catch(debugError);
+
+    this.#controller.abort();
+
+    // Repeat the last frame for the remaining frames.
+    const [buffer, timestamp] = await this.#lastFrame;
+    await Promise.all(
+      Array<Buffer>(
+        Math.max(
+          1,
+          Math.round((DEFAULT_FPS * (performance.now() - timestamp)) / 1000)
+        )
+      )
+        .fill(buffer)
+        .map(this.#writeFrame.bind(this))
+    );
+
+    // Close stdin to notify FFmpeg we are done.
+    this.#process.stdin.end();
+    await new Promise(resolve => {
+      this.#process.once('close', resolve);
+    });
+  }
+
+  /**
+   * @internal
+   */
+  async [asyncDisposeSymbol](): Promise<void> {
+    await this.stop();
+  }
+}