[Screenshotting] Fix potential race condition when screenshotting

x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/args.ts

@@ -7,7 +7,7 @@
 
 import type { ConfigType } from '../../../config';
 
-interface Viewport {
+interface WindowSize {
   height: number;
   width: number;
 }
@@ -16,12 +16,17 @@ type Proxy = ConfigType['browser']['chromium']['proxy'];
 
 interface LaunchArgs {
   userDataDir: string;
-  viewport?: Viewport;
+  windowSize?: WindowSize;
   disableSandbox?: boolean;
   proxy: Proxy;
 }
 
-export const args = ({ userDataDir, disableSandbox, viewport, proxy: proxyConfig }: LaunchArgs) => {
+export const args = ({
+  userDataDir,
+  disableSandbox,
+  windowSize,
+  proxy: proxyConfig,
+}: LaunchArgs) => {
   const flags = [
     // Disable built-in Google Translate service
     '--disable-translate',
@@ -50,11 +55,11 @@ export const args = ({ userDataDir, disableSandbox, viewport, proxy: proxyConfig
     `--mainFrameClipsContent=false`,
   ];
 
-  if (viewport) {
+  if (windowSize) {
     // NOTE: setting the window size does NOT set the viewport size: viewport and window size are different.
     // The viewport may later need to be resized depending on the position of the clip area.
     // These numbers come from the job parameters, so this is a close guess.
-    flags.push(`--window-size=${Math.floor(viewport.width)},${Math.floor(viewport.height)}`);
+    flags.push(`--window-size=${Math.floor(windowSize.width)},${Math.floor(windowSize.height)}`);
   }
 
   if (proxyConfig.enabled) {

x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/index.test.ts

@@ -11,7 +11,7 @@ import { mergeMap, take } from 'rxjs/operators';
 import type { Logger } from 'src/core/server';
 import type { ScreenshotModePluginSetup } from 'src/plugins/screenshot_mode/server';
 import { ConfigType } from '../../../config';
-import { HeadlessChromiumDriverFactory } from '.';
+import { HeadlessChromiumDriverFactory, DEFAULT_VIEWPORT } from '.';
 
 jest.mock('puppeteer');
 
@@ -70,7 +70,10 @@ describe('HeadlessChromiumDriverFactory', () => {
   describe('createPage', () => {
     it('returns browser driver, unexpected process exit observable, and close callback', async () => {
       await expect(
-        factory.createPage({ openUrlTimeout: 0 }).pipe(take(1)).toPromise()
+        factory
+          .createPage({ openUrlTimeout: 0, defaultViewport: DEFAULT_VIEWPORT })
+          .pipe(take(1))
+          .toPromise()
       ).resolves.toEqual(
         expect.objectContaining({
           driver: expect.anything(),
@@ -85,7 +88,10 @@ describe('HeadlessChromiumDriverFactory', () => {
         `Puppeteer Launch mock fail.`
       );
       expect(() =>
-        factory.createPage({ openUrlTimeout: 0 }).pipe(take(1)).toPromise()
+        factory
+          .createPage({ openUrlTimeout: 0, defaultViewport: DEFAULT_VIEWPORT })
+          .pipe(take(1))
+          .toPromise()
       ).rejects.toThrowErrorMatchingInlineSnapshot(
         `"Error spawning Chromium browser! Puppeteer Launch mock fail."`
       );
@@ -94,7 +100,7 @@ describe('HeadlessChromiumDriverFactory', () => {
     describe('close behaviour', () => {
       it('does not allow close to be called on the browse more than once', async () => {
         await factory
-          .createPage({ openUrlTimeout: 0 })
+          .createPage({ openUrlTimeout: 0, defaultViewport: DEFAULT_VIEWPORT })
           .pipe(
             take(1),
             mergeMap(async ({ close }) => {

x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/index.ts

@@ -7,6 +7,7 @@
 
 import { getDataPath } from '@kbn/utils';
 import { spawn } from 'child_process';
+import _ from 'lodash';
 import del from 'del';
 import fs from 'fs';
 import { uniq } from 'lodash';
@@ -36,6 +37,12 @@ import { getMetrics, PerformanceMetrics } from './metrics';
 
 interface CreatePageOptions {
   browserTimezone?: string;
+  defaultViewport: {
+    /** Size in pixels */
+    width?: number;
+    /** Size in pixels */
+    height?: number;
+  };
   openUrlTimeout: number;
 }
 
@@ -110,15 +117,15 @@ export class HeadlessChromiumDriverFactory {
       userDataDir: this.userDataDir,
       disableSandbox: this.config.browser.chromium.disableSandbox,
       proxy: this.config.browser.chromium.proxy,
-      viewport: DEFAULT_VIEWPORT,
+      windowSize: DEFAULT_VIEWPORT, // Approximate the default viewport size
     });
   }
 
   /*
    * Return an observable to objects which will drive screenshot capture for a page
    */
   createPage(
-    { browserTimezone, openUrlTimeout }: CreatePageOptions,
+    { browserTimezone, openUrlTimeout, defaultViewport }: CreatePageOptions,
     pLogger = this.logger
   ): Rx.Observable<CreatePageResult> {
     // FIXME: 'create' is deprecated
@@ -139,6 +146,13 @@ export class HeadlessChromiumDriverFactory {
           ignoreHTTPSErrors: true,
           handleSIGHUP: false,
           args: chromiumArgs,
+
+          // We optionally set this at page creation to reduce the chances of
+          // browser reflow. In most cases only the height needs to be adjusted
+          // before taking a screenshot.
+          // NOTE: _.defaults assigns to the target object, so we copy it.
+          // NOTE NOTE: _.defaults is not the same as { ...DEFAULT_VIEWPORT, ...defaultViewport }
+          defaultViewport: _.defaults({ ...defaultViewport }, DEFAULT_VIEWPORT),
           env: {
             TZ: browserTimezone,
           },

x-pack/plugins/screenshotting/server/layouts/create_layout.ts

@@ -5,20 +5,28 @@
  * 2.0.
  */
 
+import { map as mapRecord } from 'fp-ts/lib/Record';
 import type { LayoutParams } from '../../common/layout';
 import { LayoutTypes } from '../../common';
 import type { Layout } from '.';
 import { CanvasLayout } from './canvas_layout';
 import { PreserveLayout } from './preserve_layout';
 import { PrintLayout } from './print_layout';
 
+/**
+ * We naively round all numeric values in the object, this will break screenshotting
+ * if ever a have a non-number set as a value, but this points to an issue
+ * in the code responsible for creating the dimensions object.
+ */
+const roundNumbers = mapRecord(Math.round);
+
 export function createLayout({ id, dimensions, selectors, ...config }: LayoutParams): Layout {
   if (dimensions && id === LayoutTypes.PRESERVE_LAYOUT) {
-    return new PreserveLayout(dimensions, selectors);
+    return new PreserveLayout(roundNumbers(dimensions), selectors);
   }
 
   if (dimensions && id === LayoutTypes.CANVAS) {
-    return new CanvasLayout(dimensions);
+    return new CanvasLayout(roundNumbers(dimensions));
   }
 
   // layoutParams is optional as PrintLayout doesn't use it

x-pack/plugins/screenshotting/server/screenshots/index.test.ts

@@ -376,7 +376,7 @@ describe('Screenshot Observable Pipeline', () => {
                     "height": 1200,
                     "left": 0,
                     "top": 0,
-                    "width": 1800,
+                    "width": 1950,
                   },
                   "scroll": Object {
                     "x": 0,

x-pack/plugins/screenshotting/server/screenshots/index.ts

@@ -68,38 +68,47 @@ export function getScreenshots(
     timeouts: { openUrl: openUrlTimeout },
   } = options;
 
-  return browserDriverFactory.createPage({ browserTimezone, openUrlTimeout }, logger).pipe(
-    mergeMap(({ driver, unexpectedExit$, metrics$, close }) => {
-      apmCreatePage?.end();
-      metrics$.subscribe(({ cpu, memory }) => {
-        apmTrans?.setLabel('cpu', cpu, false);
-        apmTrans?.setLabel('memory', memory, false);
-      });
-      unexpectedExit$.subscribe({ error: () => apmTrans?.end() });
+  return browserDriverFactory
+    .createPage(
+      {
+        browserTimezone,
+        openUrlTimeout,
+        defaultViewport: { height: layout.height, width: layout.width },
+      },
+      logger
+    )
+    .pipe(
+      mergeMap(({ driver, unexpectedExit$, metrics$, close }) => {
+        apmCreatePage?.end();
+        metrics$.subscribe(({ cpu, memory }) => {
+          apmTrans?.setLabel('cpu', cpu, false);
+          apmTrans?.setLabel('memory', memory, false);
+        });
+        unexpectedExit$.subscribe({ error: () => apmTrans?.end() });
 
-      const screen = new ScreenshotObservableHandler(driver, logger, layout, options);
+        const screen = new ScreenshotObservableHandler(driver, logger, layout, options);
 
-      return from(options.urls).pipe(
-        concatMap((url, index) =>
-          screen.setupPage(index, url, apmTrans).pipe(
-            catchError((error) => {
-              screen.checkPageIsOpen(); // this fails the job if the browser has closed
+        return from(options.urls).pipe(
+          concatMap((url, index) =>
+            screen.setupPage(index, url, apmTrans).pipe(
+              catchError((error) => {
+                screen.checkPageIsOpen(); // this fails the job if the browser has closed
 
-              logger.error(error);
-              return of({ ...DEFAULT_SETUP_RESULT, error }); // allow failover screenshot capture
-            }),
-            takeUntil(unexpectedExit$),
-            screen.getScreenshots()
-          )
-        ),
-        take(options.urls.length),
-        toArray(),
-        mergeMap((results) => {
-          // At this point we no longer need the page, close it.
-          return close().pipe(mapTo({ layout, metrics$, results }));
-        })
-      );
-    }),
-    first()
-  );
+                logger.error(error);
+                return of({ ...DEFAULT_SETUP_RESULT, error }); // allow failover screenshot capture
+              }),
+              takeUntil(unexpectedExit$),
+              screen.getScreenshots()
+            )
+          ),
+          take(options.urls.length),
+          toArray(),
+          mergeMap((results) => {
+            // At this point we no longer need the page, close it.
+            return close().pipe(mapTo({ layout, metrics$, results }));
+          })
+        );
+      }),
+      first()
+    );
 }

x-pack/plugins/screenshotting/server/screenshots/observable.ts

@@ -11,7 +11,7 @@ import { catchError, mergeMap, switchMapTo, timeoutWith } from 'rxjs/operators';
 import type { Logger } from 'src/core/server';
 import type { Layout as ScreenshotModeLayout } from 'src/plugins/screenshot_mode/common';
 import type { ConditionalHeaders, HeadlessChromiumDriver } from '../browsers';
-import { getChromiumDisconnectedError } from '../browsers';
+import { getChromiumDisconnectedError, DEFAULT_VIEWPORT } from '../browsers';
 import type { Layout } from '../layouts';
 import type { ElementsPositionAndAttribute } from './get_element_position_data';
 import { getElementPositionAndAttributes } from './get_element_position_data';
@@ -107,12 +107,9 @@ interface PageSetupResults {
   error?: Error;
 }
 
-const DEFAULT_SCREENSHOT_CLIP_HEIGHT = 1200;
-const DEFAULT_SCREENSHOT_CLIP_WIDTH = 1800;
-
 const getDefaultElementPosition = (dimensions: { height?: number; width?: number } | null) => {
-  const height = dimensions?.height || DEFAULT_SCREENSHOT_CLIP_HEIGHT;
-  const width = dimensions?.width || DEFAULT_SCREENSHOT_CLIP_WIDTH;
+  const height = dimensions?.height || DEFAULT_VIEWPORT.height;
+  const width = dimensions?.width || DEFAULT_VIEWPORT.width;
 
   return [
     {
@@ -130,8 +127,7 @@ const getDefaultElementPosition = (dimensions: { height?: number; width?: number
  * provided by the browser.
  */
 const getDefaultViewPort = () => ({
-  height: DEFAULT_SCREENSHOT_CLIP_HEIGHT,
-  width: DEFAULT_SCREENSHOT_CLIP_WIDTH,
+  ...DEFAULT_VIEWPORT,
   zoom: 1,
 });
 
@@ -180,14 +176,14 @@ export class ScreenshotObservableHandler {
     const waitTimeout = this.options.timeouts.waitForElements;
 
     return defer(() => getNumberOfItems(driver, this.logger, waitTimeout, this.layout)).pipe(
-      mergeMap((itemsCount) => {
-        // set the viewport to the dimentions from the job, to allow elements to flow into the expected layout
+      mergeMap(async (itemsCount) => {
+        // set the viewport to the dimensions from the job, to allow elements to flow into the expected layout
         const viewport = this.layout.getViewport(itemsCount) || getDefaultViewPort();
 
-        return forkJoin([
-          driver.setViewport(viewport, this.logger),
-          waitForVisualizations(driver, this.logger, waitTimeout, itemsCount, this.layout),
-        ]);
+        // Set the viewport allowing time for the browser to handle reflow and redraw
+        // before checking for readiness of visualizations.
+        await driver.setViewport(viewport, this.logger);
+        await waitForVisualizations(driver, this.logger, waitTimeout, itemsCount, this.layout);
       }),
       this.waitUntil(waitTimeout, 'wait for elements')
     );