feat(trace): streaming snapshots (#5133)

- Instead of capturing snapshots on demand, we now stream them
  from each frame every 100ms.
- Certain actions can also force snapshots at particular moment using
  "checkpoints".
- Trace viewer is able to show the page snapshot at a particular
  timestamp, or using a "checkpoint" snapshot.
- Small optimization to not process stylesheets if CSSOM was not used.
  There still is a lot of room for improvement.

src/cli/cli.ts

@@ -202,6 +202,8 @@ async function launchContext(options: Options, headless: boolean): Promise<{ bro
   if (contextOptions.isMobile && browserType.name() === 'firefox')
     contextOptions.isMobile = undefined;
 
+  if (process.env.PWTRACE)
+    (contextOptions as any)._traceDir = path.join(process.cwd(), '.trace');
 
   // Proxy
 

src/cli/traceViewer/screenshotGenerator.ts

@@ -19,8 +19,7 @@ import * as path from 'path';
 import * as playwright from '../../..';
 import * as util from 'util';
 import { SnapshotRouter } from './snapshotRouter';
-import { actionById, ActionEntry, ContextEntry, TraceModel } from './traceModel';
-import type { PageSnapshot } from '../../trace/traceTypes';
+import { actionById, ActionEntry, ContextEntry, PageEntry, TraceModel } from './traceModel';
 
 const fsReadFileAsync = util.promisify(fs.readFile.bind(fs));
 const fsWriteFileAsync = util.promisify(fs.writeFile.bind(fs));
@@ -39,20 +38,18 @@ export class ScreenshotGenerator {
   }
 
   generateScreenshot(actionId: string): Promise<Buffer | undefined> {
-    const { context, action } = actionById(this._traceModel, actionId);
-    if (!action.action.snapshot)
-      return Promise.resolve(undefined);
+    const { context, action, page } = actionById(this._traceModel, actionId);
     if (!this._rendering.has(action)) {
-      this._rendering.set(action, this._render(context, action).then(body => {
+      this._rendering.set(action, this._render(context, page, action).then(body => {
         this._rendering.delete(action);
         return body;
       }));
     }
     return this._rendering.get(action)!;
   }
 
-  private async _render(contextEntry: ContextEntry, actionEntry: ActionEntry): Promise<Buffer | undefined> {
-    const imageFileName = path.join(this._traceStorageDir, actionEntry.action.snapshot!.sha1 + '-screenshot.png');
+  private async _render(contextEntry: ContextEntry, pageEntry: PageEntry, actionEntry: ActionEntry): Promise<Buffer | undefined> {
+    const imageFileName = path.join(this._traceStorageDir, actionEntry.action.timestamp + '-screenshot.png');
     try {
       return await fsReadFileAsync(imageFileName);
     } catch (e) {
@@ -70,27 +67,24 @@ export class ScreenshotGenerator {
     });
 
     try {
-      const snapshotPath = path.join(this._traceStorageDir, action.snapshot!.sha1);
-      let snapshot;
-      try {
-        snapshot = await fsReadFileAsync(snapshotPath, 'utf8');
-      } catch (e) {
-        console.log(`Unable to read snapshot at ${snapshotPath}`); // eslint-disable-line no-console
-        return;
-      }
-      const snapshotObject = JSON.parse(snapshot) as PageSnapshot;
       const snapshotRouter = new SnapshotRouter(this._traceStorageDir);
-      snapshotRouter.selectSnapshot(snapshotObject, contextEntry);
+      const snapshots = action.snapshots || [];
+      const snapshotId = snapshots.length ? snapshots[0].snapshotId : undefined;
+      const snapshotTimestamp = action.startTime;
+      const pageUrl = await snapshotRouter.selectSnapshot(contextEntry, pageEntry, snapshotId, snapshotTimestamp);
       page.route('**/*', route => snapshotRouter.route(route));
-      const url = snapshotObject.frames[0].url;
-      console.log('Generating screenshot for ' + action.action, snapshotObject.frames[0].url); // eslint-disable-line no-console
-      await page.goto(url);
+      console.log('Generating screenshot for ' + action.action, pageUrl); // eslint-disable-line no-console
+      await page.goto(pageUrl);
 
-      const element = await page.$(action.selector || '*[__playwright_target__]');
-      if (element) {
-        await element.evaluate(e => {
-          e.style.backgroundColor = '#ff69b460';
-        });
+      try {
+        const element = await page.$(action.selector || '*[__playwright_target__]');
+        if (element) {
+          await element.evaluate(e => {
+            e.style.backgroundColor = '#ff69b460';
+          });
+        }
+      } catch (e) {
+        console.log(e); // eslint-disable-line no-console
       }
       const imageData = await page.screenshot();
       await fsWriteFileAsync(imageFileName, imageData);

src/cli/traceViewer/snapshotRouter.ts

@@ -17,63 +17,125 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import * as util from 'util';
-import type { Route } from '../../..';
+import type { Frame, Route } from '../../..';
 import { parsedURL } from '../../client/clientHelper';
-import type { FrameSnapshot, NetworkResourceTraceEvent, PageSnapshot } from '../../trace/traceTypes';
-import { ContextEntry } from './traceModel';
+import { ContextEntry, PageEntry, trace } from './traceModel';
 
 const fsReadFileAsync = util.promisify(fs.readFile.bind(fs));
 
 export class SnapshotRouter {
   private _contextEntry: ContextEntry | undefined;
   private _unknownUrls = new Set<string>();
-  private _traceStorageDir: string;
-  private _frameBySrc = new Map<string, FrameSnapshot>();
+  private _resourcesDir: string;
+  private _snapshotFrameIdToSnapshot = new Map<string, trace.FrameSnapshot>();
+  private _pageUrl = '';
+  private _frameToSnapshotFrameId = new Map<Frame, string>();
 
-  constructor(traceStorageDir: string) {
-    this._traceStorageDir = traceStorageDir;
+  constructor(resourcesDir: string) {
+    this._resourcesDir = resourcesDir;
   }
 
-  selectSnapshot(snapshot: PageSnapshot, contextEntry: ContextEntry) {
-    this._frameBySrc.clear();
+  // Returns the url to navigate to.
+  async selectSnapshot(contextEntry: ContextEntry, pageEntry: PageEntry, snapshotId?: string, timestamp?: number): Promise<string> {
     this._contextEntry = contextEntry;
-    for (const frameSnapshot of snapshot.frames)
-      this._frameBySrc.set(frameSnapshot.url, frameSnapshot);
+    if (!snapshotId && !timestamp)
+      return 'data:text/html,Snapshot is not available';
+
+    const lastSnapshotEvent = new Map<string, trace.FrameSnapshotTraceEvent>();
+    for (const [frameId, snapshots] of pageEntry.snapshotsByFrameId) {
+      for (const snapshot of snapshots) {
+        const current = lastSnapshotEvent.get(frameId);
+        // Prefer snapshot with exact id.
+        const exactMatch = snapshotId && snapshot.snapshotId === snapshotId;
+        const currentExactMatch = current && snapshotId && current.snapshotId === snapshotId;
+        // If not available, prefer the latest snapshot before the timestamp.
+        const timestampMatch = timestamp && snapshot.timestamp <= timestamp;
+        if (exactMatch || (timestampMatch && !currentExactMatch))
+          lastSnapshotEvent.set(frameId, snapshot);
+      }
+    }
+
+    this._snapshotFrameIdToSnapshot.clear();
+    for (const [frameId, event] of lastSnapshotEvent) {
+      const buffer = await this._readSha1(event.sha1);
+      if (!buffer)
+        continue;
+      try {
+        const snapshot = JSON.parse(buffer.toString('utf8')) as trace.FrameSnapshot;
+        // Request url could come lower case, so we always normalize to lower case.
+        this._snapshotFrameIdToSnapshot.set(frameId.toLowerCase(), snapshot);
+      } catch (e) {
+      }
+    }
+
+    const mainFrameSnapshot = lastSnapshotEvent.get('');
+    if (!mainFrameSnapshot)
+      return 'data:text/html,Snapshot is not available';
+
+    if (!mainFrameSnapshot.frameUrl.startsWith('http'))
+      this._pageUrl = 'http://playwright.snapshot/';
+    else
+      this._pageUrl = mainFrameSnapshot.frameUrl;
+    return this._pageUrl;
   }
 
   async route(route: Route) {
     const url = route.request().url();
-    if (this._frameBySrc.has(url)) {
-      const frameSnapshot = this._frameBySrc.get(url)!;
+    const frame = route.request().frame();
+
+    if (route.request().isNavigationRequest()) {
+      let snapshotFrameId: string | undefined;
+      if (url === this._pageUrl) {
+        snapshotFrameId = '';
+      } else {
+        snapshotFrameId = url.substring(url.indexOf('://') + 3);
+        if (snapshotFrameId.endsWith('/'))
+          snapshotFrameId = snapshotFrameId.substring(0, snapshotFrameId.length - 1);
+        // Request url could come lower case, so we always normalize to lower case.
+        snapshotFrameId = snapshotFrameId.toLowerCase();
+      }
+
+      const snapshot = this._snapshotFrameIdToSnapshot.get(snapshotFrameId);
+      if (!snapshot) {
+        route.fulfill({
+          contentType: 'text/html',
+          body: 'data:text/html,Snapshot is not available',
+        });
+        return;
+      }
+
+      this._frameToSnapshotFrameId.set(frame, snapshotFrameId);
       route.fulfill({
         contentType: 'text/html',
-        body: Buffer.from(frameSnapshot.html),
+        body: snapshot.html,
       });
       return;
     }
 
-    const frameSrc = route.request().frame().url();
-    const frameSnapshot = this._frameBySrc.get(frameSrc);
-    if (!frameSnapshot)
+    const snapshotFrameId = this._frameToSnapshotFrameId.get(frame);
+    if (snapshotFrameId === undefined)
+      return this._routeUnknown(route);
+    const snapshot = this._snapshotFrameIdToSnapshot.get(snapshotFrameId);
+    if (!snapshot)
       return this._routeUnknown(route);
 
     // Find a matching resource from the same context, preferrably from the same frame.
     // Note: resources are stored without hash, but page may reference them with hash.
-    let resource: NetworkResourceTraceEvent | null = null;
+    let resource: trace.NetworkResourceTraceEvent | null = null;
     const resourcesWithUrl = this._contextEntry!.resourcesByUrl.get(removeHash(url)) || [];
     for (const resourceEvent of resourcesWithUrl) {
-      if (resource && resourceEvent.frameId !== frameSnapshot.frameId)
+      if (resource && resourceEvent.frameId !== snapshotFrameId)
         continue;
       resource = resourceEvent;
-      if (resourceEvent.frameId === frameSnapshot.frameId)
+      if (resourceEvent.frameId === snapshotFrameId)
         break;
     }
     if (!resource)
       return this._routeUnknown(route);
 
     // This particular frame might have a resource content override, for example when
     // stylesheet is modified using CSSOM.
-    const resourceOverride = frameSnapshot.resourceOverrides.find(o => o.url === url);
+    const resourceOverride = snapshot.resourceOverrides.find(o => o.url === url);
     const overrideSha1 = resourceOverride ? resourceOverride.sha1 : undefined;
     const resourceData = await this._readResource(resource, overrideSha1);
     if (!resourceData)
@@ -98,18 +160,24 @@ export class SnapshotRouter {
     route.abort();
   }
 
-  private async _readResource(event: NetworkResourceTraceEvent, overrideSha1: string | undefined) {
+  private async _readSha1(sha1: string) {
     try {
-      const body = await fsReadFileAsync(path.join(this._traceStorageDir, overrideSha1 || event.sha1));
-      return {
-        contentType: event.contentType,
-        body,
-        headers: event.responseHeaders,
-      };
+      return await fsReadFileAsync(path.join(this._resourcesDir, sha1));
     } catch (e) {
       return undefined;
     }
   }
+
+  private async _readResource(event: trace.NetworkResourceTraceEvent, overrideSha1: string | undefined) {
+    const body = await this._readSha1(overrideSha1 || event.sha1);
+    if (!body)
+      return;
+    return {
+      contentType: event.contentType,
+      body,
+      headers: event.responseHeaders,
+    };
+  }
 }
 
 function removeHash(url: string) {

src/cli/traceViewer/traceModel.ts

@@ -46,6 +46,7 @@ export type PageEntry = {
   actions: ActionEntry[];
   interestingEvents: InterestingPageEvent[];
   resources: trace.NetworkResourceTraceEvent[];
+  snapshotsByFrameId: Map<string, trace.FrameSnapshotTraceEvent[]>;
 }
 
 export type ActionEntry = {
@@ -93,6 +94,7 @@ export function readTraceFile(events: trace.TraceEvent[], traceModel: TraceModel
           actions: [],
           resources: [],
           interestingEvents: [],
+          snapshotsByFrameId: new Map(),
         };
         pageEntries.set(event.pageId, pageEntry);
         contextEntries.get(event.contextId)!.pages.push(pageEntry);
@@ -144,6 +146,13 @@ export function readTraceFile(events: trace.TraceEvent[], traceModel: TraceModel
         pageEntry.interestingEvents.push(event);
         break;
       }
+      case 'snapshot': {
+        const pageEntry = pageEntries.get(event.pageId!)!;
+        if (!pageEntry.snapshotsByFrameId.has(event.frameId))
+          pageEntry.snapshotsByFrameId.set(event.frameId, []);
+        pageEntry.snapshotsByFrameId.get(event.frameId)!.push(event);
+        break;
+      }
     }
 
     const contextEntry = contextEntries.get(event.contextId)!;

src/cli/traceViewer/traceViewer.ts

@@ -21,7 +21,7 @@ import * as util from 'util';
 import { ScreenshotGenerator } from './screenshotGenerator';
 import { SnapshotRouter } from './snapshotRouter';
 import { readTraceFile, TraceModel } from './traceModel';
-import type { ActionTraceEvent, PageSnapshot, TraceEvent } from '../../trace/traceTypes';
+import type { ActionTraceEvent, TraceEvent } from '../../trace/traceTypes';
 
 const fsReadFileAsync = util.promisify(fs.readFile.bind(fs));
 
@@ -92,25 +92,20 @@ class TraceViewer {
     await uiPage.exposeBinding('readFile', async (_, path: string) => {
       return fs.readFileSync(path).toString();
     });
-    await uiPage.exposeBinding('renderSnapshot', async (_, action: ActionTraceEvent) => {
+    await uiPage.exposeBinding('renderSnapshot', async (_, arg: { action: ActionTraceEvent, snapshot: { name: string, snapshotId?: string } }) => {
+      const { action, snapshot } = arg;
       if (!this._document)
         return;
       try {
-        if (!action.snapshot) {
-          const snapshotFrame = uiPage.frames()[1];
-          await snapshotFrame.goto('data:text/html,No snapshot available');
-          return;
-        }
-
-        const snapshot = await fsReadFileAsync(path.join(this._document.resourcesDir, action.snapshot!.sha1), 'utf8');
-        const snapshotObject = JSON.parse(snapshot) as PageSnapshot;
         const contextEntry = this._document.model.contexts.find(entry => entry.created.contextId === action.contextId)!;
-        this._document.snapshotRouter.selectSnapshot(snapshotObject, contextEntry);
+        const pageEntry = contextEntry.pages.find(entry => entry.created.pageId === action.pageId)!;
+        const snapshotTime = snapshot.name === 'before' ? action.startTime : (snapshot.name === 'after' ? action.endTime : undefined);
+        const pageUrl = await this._document.snapshotRouter.selectSnapshot(contextEntry, pageEntry, snapshot.snapshotId, snapshotTime);
 
         // TODO: fix Playwright bug where frame.name is lost (empty).
         const snapshotFrame = uiPage.frames()[1];
         try {
-          await snapshotFrame.goto(snapshotObject.frames[0].url);
+          await snapshotFrame.goto(pageUrl);
         } catch (e) {
           if (!e.message.includes('frame was detached'))
             console.error(e);

src/cli/traceViewer/web/index.tsx

@@ -25,7 +25,7 @@ declare global {
   interface Window {
     getTraceModel(): Promise<TraceModel>;
     readFile(filePath: string): Promise<string>;
-    renderSnapshot(action: trace.ActionTraceEvent): void;
+    renderSnapshot(arg: { action: trace.ActionTraceEvent, snapshot: { name: string, snapshotId?: string } }): void;
   }
 }
 

src/cli/traceViewer/web/ui/propertiesTabbedPane.css

@@ -68,6 +68,28 @@
   font-weight: 600;
 }
 
+.snapshot-tab {
+  display: flex;
+  flex-direction: column;
+  align-items: stretch;
+}
+
+.snapshot-controls {
+  flex: 0 0 24px;
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+}
+
+.snapshot-toggle {
+  padding: 5px 10px;
+  cursor: pointer;
+}
+
+.snapshot-toggle.toggled {
+  background: var(--inactive-focus-ring);
+}
+
 .snapshot-wrapper {
   flex: auto;
   margin: 1px;