Merge pull request #791 from Mike-Dax/tracing

Record CPU traces from screenshot runs

README.md

@@ -361,6 +361,7 @@ Options:
   -V, --viewport                   Viewport.                                              [array] [default: ["800x600"]]
       --disableCssAnimation        Disable CSS animation and transition.                       [boolean] [default: true]
       --disableWaitAssets          Disable waiting for requested assets                       [boolean] [default: false]
+      --trace                      Emit Chromium trace files per screenshot.                  [boolean] [default: false]
       --silent                                                                                [boolean] [default: false]
       --verbose                                                                               [boolean] [default: false]
       --forwardConsoleLogs         Forward in-page console logs to the user's console.        [boolean] [default: false]

packages/storycap/src/node/capturing-browser.ts

@@ -22,6 +22,7 @@ import {
   InvalidVariantKeysReason,
 } from '../shared/screenshot-options-helper';
 import { Logger } from './logger';
+import { FileSystem } from './file';
 
 /**
  *
@@ -349,6 +350,9 @@ export class CapturingBrowser extends StoryPreviewBrowser {
    * @param requestId - Represents an identifier for the screenshot
    * @param variantKey - Variant identifier for the screenshot
    * @param retryCount - The number which represents how many attempting to capture this story and variant
+   * @param logger - Logger instance
+   * @param forwardConsoleLogs - Whether to forward logs from the page to the user's console
+   * @param trace - Whether to record a CPU trace per screenshot
    *
    * @returns PNG buffer, whether the capturing process is succeeded or not, additional variant keys if they are emitted, and file name suffix for default the default variant.
    *
@@ -364,6 +368,8 @@ export class CapturingBrowser extends StoryPreviewBrowser {
     retryCount: number,
     logger: Logger,
     forwardConsoleLogs: boolean,
+    trace: boolean,
+    fileSystem: FileSystem,
   ): Promise<ScreenshotResult> {
     this.currentRequestId = requestId;
     this.currentVariantKey = variantKey;
@@ -393,6 +399,14 @@ export class CapturingBrowser extends StoryPreviewBrowser {
 
     this.page.on('console', onConsoleLog);
 
+    if (trace) {
+      // Begin CPU trace, don't write to file, store it in memory
+      await this.page.tracing.start();
+    }
+
+    // Capture this outside so it can be used for the filePath generation for the trace
+    let defaultVariantSuffix: string | undefined;
+
     try {
       await this.setCurrentStory(story, { forceRerender: true });
 
@@ -413,6 +427,9 @@ export class CapturingBrowser extends StoryPreviewBrowser {
         emittedScreenshotOptions = pickupWithVariantKey(this.baseScreenshotOptions, this.currentVariantKey);
       }
 
+      // Set defaultVariantSuffix as soon as it's known
+      defaultVariantSuffix = emittedScreenshotOptions.defaultVariantSuffix;
+
       const mergedScreenshotOptions = mergeScreenshotOptions(this.baseScreenshotOptions, emittedScreenshotOptions);
 
       // Get keys for variants included in the screenshot options in order to queue capturing them after this sequence.
@@ -467,10 +484,19 @@ export class CapturingBrowser extends StoryPreviewBrowser {
         buffer,
         succeeded: true,
         variantKeysToPush,
-        defaultVariantSuffix: emittedScreenshotOptions.defaultVariantSuffix,
+        defaultVariantSuffix,
       };
     } finally {
       this.page.off('console', onConsoleLog);
+
+      if (trace) {
+        // Finish CPU trace.
+        const traceBuffer = await this.page.tracing.stop();
+
+        // Calculate the suffix and save the trace to the file.
+        const suffix = variantKey.isDefault && defaultVariantSuffix ? [defaultVariantSuffix] : variantKey.keys;
+        await fileSystem.saveTrace(story.kind, story.story, suffix, traceBuffer);
+      }
     }
   }
 }

packages/storycap/src/node/cli.ts

@@ -34,6 +34,7 @@ function createOptions(): MainOptions {
       default: false,
       description: 'Disable waiting for requested assets',
     })
+    .option('trace', { boolean: true, default: false, description: 'Emit Chromium trace files per screenshot.' })
     .option('silent', { boolean: true, default: false })
     .option('verbose', { boolean: true, default: false })
     .option('forwardConsoleLogs', {
@@ -131,6 +132,7 @@ function createOptions(): MainOptions {
     stateChangeDelay,
     disableCssAnimation,
     disableWaitAssets,
+    trace,
     listDevices,
     chromiumChannel,
     chromiumPath,
@@ -185,6 +187,7 @@ function createOptions(): MainOptions {
     stateChangeDelay,
     disableCssAnimation,
     disableWaitAssets,
+    trace,
     forwardConsoleLogs,
     chromiumChannel: chromiumChannel as ChromeChannel,
     chromiumPath,

packages/storycap/src/node/file.ts

@@ -6,6 +6,20 @@ import sanitize from 'sanitize-filename';
 export class FileSystem {
   constructor(private opt: MainOptions) {}
 
+  private getPath(kind: string, story: string, suffix: string[], extension: string) {
+    const name = this.opt.flat
+      ? sanitize((kind + '_' + story).replace(/\//g, '_'))
+      : kind
+          .split('/')
+          .map(k => sanitize(k))
+          .join('/') +
+        '/' +
+        sanitize(story);
+    const filePath = path.join(this.opt.outDir, name + (suffix.length ? `_${suffix.join('_')}` : '') + extension);
+
+    return filePath;
+  }
+
   /**
    *
    * Save captured buffer as a PNG image.
@@ -17,16 +31,28 @@ export class FileSystem {
    * @returns Absolute file path
    *
    **/
-  async save(kind: string, story: string, suffix: string[], buffer: Buffer) {
-    const name = this.opt.flat
-      ? sanitize((kind + '_' + story).replace(/\//g, '_'))
-      : kind
-          .split('/')
-          .map(k => sanitize(k))
-          .join('/') +
-        '/' +
-        sanitize(story);
-    const filePath = path.join(this.opt.outDir, name + (suffix.length ? `_${suffix.join('_')}` : '') + '.png');
+  async saveScreenshot(kind: string, story: string, suffix: string[], buffer: Buffer) {
+    const filePath = this.getPath(kind, story, suffix, '.png');
+
+    await fs.mkdir(path.dirname(filePath), { recursive: true });
+    await fs.writeFile(filePath, buffer);
+
+    return filePath;
+  }
+
+  /**
+   *
+   * Save captured tracing buffer as a json file.
+   *
+   * @param kind - Story kind
+   * @param story - Name of this story
+   * @param suffix - File name suffix
+   * @param buffer - PNG image buffer to save
+   * @returns Absolute file path
+   *
+   **/
+  async saveTrace(kind: string, story: string, suffix: string[], buffer: Buffer) {
+    const filePath = this.getPath(kind, story, [...suffix, 'trace'], '.json');
 
     await fs.mkdir(path.dirname(filePath), { recursive: true });
     await fs.writeFile(filePath, buffer);

packages/storycap/src/node/main.ts

@@ -99,6 +99,7 @@ export async function main(mainOptions: MainOptions) {
       fileSystem,
       logger,
       forwardConsoleLogs: mainOptions.forwardConsoleLogs,
+      trace: mainOptions.trace,
     }).execute();
     logger.debug('Ended ScreenshotService execution.');
     return captured;

packages/storycap/src/node/screenshot-service.ts

@@ -49,6 +49,7 @@ export type ScreenshotServiceOptions = {
   fileSystem: FileSystem;
   stories: Story[];
   forwardConsoleLogs: boolean;
+  trace: boolean;
 };
 
 /**
@@ -65,6 +66,7 @@ export function createScreenshotService({
   stories,
   workers,
   forwardConsoleLogs,
+  trace,
 }: ScreenshotServiceOptions): ScreenshotService {
   const service = createExecutionService(
     workers,
@@ -73,7 +75,7 @@ export function createScreenshotService({
       async worker => {
         // Delegate the request to the worker.
         const [result, elapsedTime] = await time(
-          worker.screenshot(rid, story, variantKey, count, logger, forwardConsoleLogs),
+          worker.screenshot(rid, story, variantKey, count, logger, forwardConsoleLogs, trace, fileSystem),
         );
 
         const { succeeded, buffer, variantKeysToPush, defaultVariantSuffix } = result;
@@ -87,7 +89,7 @@ export function createScreenshotService({
 
         if (buffer) {
           const suffix = variantKey.isDefault && defaultVariantSuffix ? [defaultVariantSuffix] : variantKey.keys;
-          const path = await fileSystem.save(story.kind, story.story, suffix, buffer);
+          const path = await fileSystem.saveScreenshot(story.kind, story.story, suffix, buffer);
           logger.log(`Screenshot stored: ${logger.color.magenta(path)} in ${elapsedTime} msec.`);
           return true;
         }

packages/storycap/src/node/types.ts

@@ -40,6 +40,7 @@ export interface MainOptions extends BaseBrowserOptions {
   exclude: string[];
   disableCssAnimation: boolean;
   disableWaitAssets: boolean;
+  trace: boolean;
   forwardConsoleLogs: boolean;
   parallel: number;
   shard: ShardOptions;

packages/storycap/src/shared/screenshot-options-helper.ts

@@ -14,6 +14,7 @@ const defaultScreenshotOptions = {
   captureBeyondViewport: true,
   clip: null,
   forwardConsoleLogs: false,
+  trace: false,
 } as const;
 
 /**