[Reporting] Close chromium before compiling pdf (#120223)

* added "close" method to the factory output so that consumers can close the browser once they are done with the page * update test reference to exit$ * removed old files * fix variable name * remove unused import and fix typo in mock * added comments and documentation for close functionality * updated implemenatation of close to not call .kill if page is already closed * added a test describing the close behaviour we want, also moved the page.isClosed() check to the private .kill function to be DRY * avoid emitting a log unnecessarily Co-authored-by: Kibana Machine <[email protected]>
elastic · Dec 13, 2021 · cbe06ea · cbe06ea
1 parent 797bc49
commit cbe06ea
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 20 deletions.
diff --git a/x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/index.test.ts b/x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/index.test.ts
@@ -7,7 +7,7 @@
 
 import puppeteer from 'puppeteer';
 import * as Rx from 'rxjs';
-import { take } from 'rxjs/operators';
+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';
@@ -27,6 +27,7 @@ describe('HeadlessChromiumDriverFactory', () => {
   let logger: jest.Mocked<Logger>;
   let screenshotMode: jest.Mocked<ScreenshotModePluginSetup>;
   let factory: HeadlessChromiumDriverFactory;
+  let mockBrowser: jest.Mocked<puppeteer.Browser>;
 
   beforeEach(async () => {
     logger = {
@@ -38,7 +39,8 @@ describe('HeadlessChromiumDriverFactory', () => {
     } as unknown as typeof logger;
     screenshotMode = {} as unknown as typeof screenshotMode;
 
-    (puppeteer as jest.Mocked<typeof puppeteer>).launch.mockResolvedValue({
+    let pageClosed = false;
+    mockBrowser = {
       newPage: jest.fn().mockResolvedValue({
         target: jest.fn(() => ({
           createCDPSession: jest.fn().mockResolvedValue({
@@ -47,10 +49,17 @@ describe('HeadlessChromiumDriverFactory', () => {
         })),
         emulateTimezone: jest.fn(),
         setDefaultTimeout: jest.fn(),
+        isClosed: jest.fn(() => {
+          return pageClosed;
+        }),
+      }),
+      close: jest.fn(() => {
+        pageClosed = true;
       }),
-      close: jest.fn(),
       process: jest.fn(),
-    } as unknown as puppeteer.Browser);
+    } as unknown as jest.Mocked<puppeteer.Browser>;
+
+    (puppeteer as jest.Mocked<typeof puppeteer>).launch.mockResolvedValue(mockBrowser);
 
     factory = new HeadlessChromiumDriverFactory(screenshotMode, config, logger, path);
     jest.spyOn(factory, 'getBrowserLogger').mockReturnValue(Rx.EMPTY);
@@ -59,13 +68,14 @@ describe('HeadlessChromiumDriverFactory', () => {
   });
 
   describe('createPage', () => {
-    it('returns browser driver and process exit observable', async () => {
+    it('returns browser driver, unexpected process exit observable, and close callback', async () => {
       await expect(
         factory.createPage({ openUrlTimeout: 0 }).pipe(take(1)).toPromise()
       ).resolves.toEqual(
         expect.objectContaining({
           driver: expect.anything(),
-          exit$: expect.anything(),
+          unexpectedExit$: expect.anything(),
+          close: expect.anything(),
         })
       );
     });
@@ -80,5 +90,24 @@ describe('HeadlessChromiumDriverFactory', () => {
         `"Error spawning Chromium browser! Puppeteer Launch mock fail."`
       );
     });
+
+    describe('close behaviour', () => {
+      it('does not allow close to be called on the browse more than once', async () => {
+        await factory
+          .createPage({ openUrlTimeout: 0 })
+          .pipe(
+            take(1),
+            mergeMap(async ({ close }) => {
+              expect(mockBrowser.close).not.toHaveBeenCalled();
+              await close().toPromise();
+              await close().toPromise();
+              expect(mockBrowser.close).toHaveBeenCalledTimes(1);
+            })
+          )
+          .toPromise();
+        // Check again, after the observable completes
+        expect(mockBrowser.close).toHaveBeenCalledTimes(1);
+      });
+    });
   });
 });
diff --git a/x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/index.ts b/x-pack/plugins/screenshotting/server/browsers/chromium/driver_factory/index.ts
@@ -33,8 +33,16 @@ interface CreatePageOptions {
 
 interface CreatePageResult {
   driver: HeadlessChromiumDriver;
-  exit$: Rx.Observable<never>;
+  unexpectedExit$: Rx.Observable<never>;
   metrics$: Rx.Observable<PerformanceMetrics>;
+  /**
+   * Close the page and the browser.
+   *
+   * @note Ensure this function gets called once all actions against the page
+   * have concluded. This ensures the browser is closed and gives the OS a chance
+   * to reclaim resources like memory.
+   */
+  close: () => Rx.Observable<void>;
 }
 
 export const DEFAULT_VIEWPORT = {
@@ -152,7 +160,8 @@ export class HeadlessChromiumDriverFactory {
       logger.debug(`Browser page driver created`);
 
       const childProcess = {
-        async kill() {
+        async kill(): Promise<void> {
+          if (page.isClosed()) return;
           try {
             if (devTools && startMetrics) {
               const endMetrics = await devTools.send('Performance.getMetrics');
@@ -171,7 +180,9 @@ export class HeadlessChromiumDriverFactory {
           }
 
           try {
+            logger.debug('Attempting to close browser...');
             await browser?.close();
+            logger.debug('Browser closed.');
           } catch (err) {
             // do not throw
             logger.error(err);
@@ -180,10 +191,10 @@ export class HeadlessChromiumDriverFactory {
       };
       const { terminate$ } = safeChildProcess(logger, childProcess);
 
-      // this is adding unsubscribe logic to our observer
-      // so that if our observer unsubscribes, we terminate our child-process
+      // Ensure that the browser is closed once the observable completes.
       observer.add(() => {
-        logger.debug(`The browser process observer has unsubscribed. Closing the browser...`);
+        if (page.isClosed()) return; // avoid emitting a log unnecessarily
+        logger.debug(`It looks like the browser is no longer being used. Closing the browser...`);
         childProcess.kill(); // ignore async
       });
 
@@ -207,9 +218,14 @@ export class HeadlessChromiumDriverFactory {
       const driver = new HeadlessChromiumDriver(this.screenshotMode, this.config, page);
 
       // Rx.Observable<never>: stream to interrupt page capture
-      const exit$ = this.getPageExit(browser, page);
+      const unexpectedExit$ = this.getPageExit(browser, page);
 
-      observer.next({ driver, exit$, metrics$: metrics$.asObservable() });
+      observer.next({
+        driver,
+        unexpectedExit$,
+        metrics$: metrics$.asObservable(),
+        close: () => Rx.from(childProcess.kill()),
+      });
 
       // unsubscribe logic makes a best-effort attempt to delete the user data directory used by chromium
       observer.add(() => {

diff --git a/x-pack/plugins/screenshotting/server/browsers/mock.ts b/x-pack/plugins/screenshotting/server/browsers/mock.ts
@@ -88,7 +88,12 @@ export function createMockBrowserDriverFactory(
 ): jest.Mocked<HeadlessChromiumDriverFactory> {
   return {
     createPage: jest.fn(() =>
-      of({ driver: driver ?? createMockBrowserDriver(), exit$: NEVER, metrics$: NEVER })
+      of({
+        driver: driver ?? createMockBrowserDriver(),
+        unexpectedExit$: NEVER,
+        metrics$: NEVER,
+        close: () => of(undefined),
+      })
     ),
     diagnose: jest.fn(() => of('message')),
   } as unknown as ReturnType<typeof createMockBrowserDriverFactory>;

diff --git a/x-pack/plugins/screenshotting/server/screenshots/index.test.ts b/x-pack/plugins/screenshotting/server/screenshots/index.test.ts
@@ -343,7 +343,12 @@ describe('Screenshot Observable Pipeline', () => {
 
     it('observes page exit', async () => {
       driverFactory.createPage.mockReturnValue(
-        of({ driver, exit$: throwError('Instant timeout has fired!'), metrics$: NEVER })
+        of({
+          driver,
+          unexpectedExit$: throwError('Instant timeout has fired!'),
+          metrics$: NEVER,
+          close: () => of(undefined),
+        })
       );
 
       await expect(

diff --git a/x-pack/plugins/screenshotting/server/screenshots/index.ts b/x-pack/plugins/screenshotting/server/screenshots/index.ts
@@ -11,7 +11,7 @@ import {
   catchError,
   concatMap,
   first,
-  map,
+  mapTo,
   mergeMap,
   take,
   takeUntil,
@@ -69,13 +69,13 @@ export function getScreenshots(
   } = options;
 
   return browserDriverFactory.createPage({ browserTimezone, openUrlTimeout }, logger).pipe(
-    mergeMap(({ driver, exit$, metrics$ }) => {
+    mergeMap(({ driver, unexpectedExit$, metrics$, close }) => {
       apmCreatePage?.end();
       metrics$.subscribe(({ cpu, memory }) => {
         apmTrans?.setLabel('cpu', cpu, false);
         apmTrans?.setLabel('memory', memory, false);
       });
-      exit$.subscribe({ error: () => apmTrans?.end() });
+      unexpectedExit$.subscribe({ error: () => apmTrans?.end() });
 
       const screen = new ScreenshotObservableHandler(driver, logger, layout, options);
 
@@ -88,13 +88,16 @@ export function getScreenshots(
               logger.error(error);
               return of({ ...DEFAULT_SETUP_RESULT, error }); // allow failover screenshot capture
             }),
-            takeUntil(exit$),
+            takeUntil(unexpectedExit$),
             screen.getScreenshots()
           )
         ),
         take(options.urls.length),
         toArray(),
-        map((results) => ({ layout, metrics$, results }))
+        mergeMap((results) => {
+          // At this point we no longer need the page, close it.
+          return close().pipe(mapTo({ layout, metrics$, results }));
+        })
       );
     }),
     first()