feat(tracing): introduce _reset() and _export() (#7974)

`tracing._export({ path })` exports current tracing state into a file
and does not require tracing to be stopped.

`tracing._reset()` resets current tracing state, but keeps resources
around so they can be referenced in the future snapshots. Does not stop.

The usage pattern is:
```js
await tracing.start({ screenshots: true, snapshots: true });
// ...
await tracing._reset();
// Do stuff, it will all be in the export below.
await tracing._export({ path });
// ...
await tracing.stop();
```

src/client/tracing.ts

@@ -32,17 +32,32 @@ export class Tracing implements api.Tracing {
     });
   }
 
+  async _reset() {
+    await this._context._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
+      return await channel.tracingReset();
+    });
+  }
+
+  async _export(options: { path: string }) {
+    await this._context._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
+      await this._doExport(channel, options.path);
+    });
+  }
+
   async stop(options: { path?: string } = {}) {
     await this._context._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
       await channel.tracingStop();
-      if (options.path) {
-        const result = await channel.tracingExport();
-        const artifact = Artifact.from(result.artifact);
-        if (this._context.browser()?._remoteType)
-          artifact._isRemote = true;
-        await artifact.saveAs(options.path);
-        await artifact.delete();
-      }
+      if (options.path)
+        await this._doExport(channel, options.path);
     });
   }
+
+  private async _doExport(channel: channels.BrowserContextChannel, path: string) {
+    const result = await channel.tracingExport();
+    const artifact = Artifact.from(result.artifact);
+    if (this._context.browser()?._remoteType)
+      artifact._isRemote = true;
+    await artifact.saveAs(path);
+    await artifact.delete();
+  }
 }

src/dispatchers/browserContextDispatcher.ts

@@ -184,6 +184,10 @@ export class BrowserContextDispatcher extends Dispatcher<BrowserContext, channel
     await this._context.tracing.start(params);
   }
 
+  async tracingReset(params: channels.BrowserContextTracingResetParams): Promise<channels.BrowserContextTracingResetResult> {
+    await this._context.tracing.reset();
+  }
+
   async tracingStop(params: channels.BrowserContextTracingStopParams): Promise<channels.BrowserContextTracingStopResult> {
     await this._context.tracing.stop();
   }

src/protocol/channels.ts

@@ -656,6 +656,7 @@ export interface BrowserContextChannel extends EventTargetChannel {
   recorderSupplementEnable(params: BrowserContextRecorderSupplementEnableParams, metadata?: Metadata): Promise<BrowserContextRecorderSupplementEnableResult>;
   newCDPSession(params: BrowserContextNewCDPSessionParams, metadata?: Metadata): Promise<BrowserContextNewCDPSessionResult>;
   tracingStart(params: BrowserContextTracingStartParams, metadata?: Metadata): Promise<BrowserContextTracingStartResult>;
+  tracingReset(params?: BrowserContextTracingResetParams, metadata?: Metadata): Promise<BrowserContextTracingResetResult>;
   tracingStop(params?: BrowserContextTracingStopParams, metadata?: Metadata): Promise<BrowserContextTracingStopResult>;
   tracingExport(params?: BrowserContextTracingExportParams, metadata?: Metadata): Promise<BrowserContextTracingExportResult>;
 }
@@ -864,6 +865,9 @@ export type BrowserContextTracingStartOptions = {
   screenshots?: boolean,
 };
 export type BrowserContextTracingStartResult = void;
+export type BrowserContextTracingResetParams = {};
+export type BrowserContextTracingResetOptions = {};
+export type BrowserContextTracingResetResult = void;
 export type BrowserContextTracingStopParams = {};
 export type BrowserContextTracingStopOptions = {};
 export type BrowserContextTracingStopResult = void;

src/protocol/protocol.yml

@@ -639,6 +639,8 @@ BrowserContext:
         snapshots: boolean?
         screenshots: boolean?
 
+    tracingReset:
+
     tracingStop:
 
     tracingExport:

src/protocol/validator.ts

@@ -413,6 +413,7 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     snapshots: tOptional(tBoolean),
     screenshots: tOptional(tBoolean),
   });
+  scheme.BrowserContextTracingResetParams = tOptional(tObject({}));
   scheme.BrowserContextTracingStopParams = tOptional(tObject({}));
   scheme.BrowserContextTracingExportParams = tOptional(tObject({}));
   scheme.PageSetDefaultNavigationTimeoutNoReplyParams = tObject({

src/server/snapshot/snapshotter.ts

@@ -62,7 +62,7 @@ export class Snapshotter {
       this._initialized = true;
       await this._initialize();
     }
-    await this._runInAllFrames(`window["${this._snapshotStreamer}"].reset()`);
+    await this.reset();
 
     // Replay resources loaded in all pages.
     for (const page of this._context.pages()) {
@@ -71,6 +71,11 @@ export class Snapshotter {
     }
   }
 
+  async reset() {
+    if (this._started)
+      await this._runInAllFrames(`window["${this._snapshotStreamer}"].reset()`);
+  }
+
   async stop() {
     this._started = false;
   }

src/server/trace/common/traceEvents.ts

@@ -50,9 +50,15 @@ export type FrameSnapshotTraceEvent = {
   snapshot: FrameSnapshot,
 };
 
+export type MarkerTraceEvent = {
+  type: 'marker',
+  resetIndex?: number,
+};
+
 export type TraceEvent =
     ContextCreatedTraceEvent |
     ScreencastFrameTraceEvent |
     ActionTraceEvent |
     ResourceSnapshotTraceEvent |
-    FrameSnapshotTraceEvent;
+    FrameSnapshotTraceEvent |
+    MarkerTraceEvent;

src/server/trace/recorder/traceSnapshotter.ts

@@ -46,6 +46,14 @@ export class TraceSnapshotter extends EventEmitter implements SnapshotterDelegat
     await this._snapshotter.start();
   }
 
+  async reset() {
+    await this._snapshotter.reset();
+  }
+
+  async checkpoint() {
+    await this._writeArtifactChain;
+  }
+
   async stop(): Promise<void> {
     await this._snapshotter.stop();
     await this._writeArtifactChain;

src/server/trace/recorder/tracing.ts

@@ -17,6 +17,7 @@
 import fs from 'fs';
 import path from 'path';
 import yazl from 'yazl';
+import readline from 'readline';
 import { EventEmitter } from 'events';
 import { calculateSha1, createGuid, mkdirIfNeeded, monotonicTime } from '../../../utils/utils';
 import { Artifact } from '../../artifact';
@@ -45,9 +46,10 @@ export class Tracing implements InstrumentationListener {
   private _context: BrowserContext;
   private _traceFile: string | undefined;
   private _resourcesDir: string;
-  private _sha1s: string[] = [];
+  private _sha1s = new Set<string>();
   private _recordingTraceEvents = false;
   private _tracesDir: string;
+  private _lastReset = 0;
 
   constructor(context: BrowserContext) {
     this._context = context;
@@ -64,6 +66,7 @@ export class Tracing implements InstrumentationListener {
     // TODO: passing the same name for two contexts makes them write into a single file
     // and conflict.
     this._traceFile = path.join(this._tracesDir, (options.name || createGuid()) + '.trace');
+    this._lastReset = 0;
 
     this._appendEventChain = mkdirIfNeeded(this._traceFile);
     const event: trace.ContextCreatedTraceEvent = {
@@ -87,6 +90,16 @@ export class Tracing implements InstrumentationListener {
       await this._snapshotter.start();
   }
 
+  async reset(): Promise<void> {
+    await this._appendTraceOperation(async () => {
+      // Reset snapshots to avoid back-references.
+      await this._snapshotter.reset();
+      this._lastReset++;
+      const markerEvent: trace.MarkerTraceEvent = { type: 'marker', resetIndex: this._lastReset };
+      await fs.promises.appendFile(this._traceFile!, JSON.stringify(markerEvent) + '\n');
+    });
+  }
+
   async stop(): Promise<void> {
     if (!this._eventListeners.length)
       return;
@@ -112,25 +125,77 @@ export class Tracing implements InstrumentationListener {
   }
 
   async export(): Promise<Artifact> {
-    if (!this._traceFile || this._recordingTraceEvents)
-      throw new Error('Must start and stop tracing before exporting');
-    const zipFile = new yazl.ZipFile();
-    const failedPromise = new Promise<Artifact>((_, reject) => (zipFile as any as EventEmitter).on('error', reject));
+    if (!this._traceFile)
+      throw new Error('Must start tracing before exporting');
+    // Chain the export operation against write operations,
+    // so that neither trace file nor sha1s change during the export.
+    return await this._appendTraceOperation(async () => {
+      await this._snapshotter.checkpoint();
 
-    const succeededPromise = new Promise<Artifact>(async fulfill => {
-      zipFile.addFile(this._traceFile!, 'trace.trace');
-      const zipFileName = this._traceFile! + '.zip';
-      for (const sha1 of this._sha1s)
-        zipFile.addFile(path.join(this._resourcesDir!, sha1), path.join('resources', sha1));
-      zipFile.end();
-      await new Promise(f => {
-        zipFile.outputStream.pipe(fs.createWriteStream(zipFileName)).on('close', f);
+      const resetIndex = this._lastReset;
+      let trace = { file: this._traceFile!, sha1s: this._sha1s };
+      // Make a filtered trace if needed.
+      if (resetIndex)
+        trace = await this._filterTrace(this._traceFile!, resetIndex);
+
+      const zipFile = new yazl.ZipFile();
+      const failedPromise = new Promise<Artifact>((_, reject) => (zipFile as any as EventEmitter).on('error', reject));
+      const succeededPromise = new Promise<Artifact>(async fulfill => {
+        zipFile.addFile(trace.file, 'trace.trace');
+        const zipFileName = trace.file + '.zip';
+        for (const sha1 of trace.sha1s)
+          zipFile.addFile(path.join(this._resourcesDir!, sha1), path.join('resources', sha1));
+        zipFile.end();
+        await new Promise(f => {
+          zipFile.outputStream.pipe(fs.createWriteStream(zipFileName)).on('close', f);
+        });
+        const artifact = new Artifact(this._context, zipFileName);
+        artifact.reportFinished();
+        fulfill(artifact);
+      });
+      return Promise.race([failedPromise, succeededPromise]).finally(async () => {
+        // Remove the filtered trace.
+        if (resetIndex)
+          await fs.promises.unlink(trace.file).catch(() => {});
       });
-      const artifact = new Artifact(this._context, zipFileName);
-      artifact.reportFinished();
-      fulfill(artifact);
     });
-    return Promise.race([failedPromise, succeededPromise]);
+  }
+
+  private async _filterTrace(traceFile: string, resetIndex: number): Promise<{ file: string, sha1s: Set<string> }> {
+    const ext = path.extname(traceFile);
+    const traceFileCopy = traceFile.substring(0, traceFile.length - ext.length) + '-copy' + resetIndex + ext;
+    const sha1s = new Set<string>();
+    await new Promise<void>((resolve, reject) => {
+      const fileStream = fs.createReadStream(traceFile, 'utf8');
+      const rl = readline.createInterface({
+        input: fileStream,
+        crlfDelay: Infinity
+      });
+      let copyChain = Promise.resolve();
+      let foundMarker = false;
+      rl.on('line', line => {
+        try {
+          const event = JSON.parse(line) as trace.TraceEvent;
+          if (event.type === 'marker' && event.resetIndex === resetIndex) {
+            foundMarker = true;
+          } else if (event.type === 'resource-snapshot' || event.type === 'context-options' || foundMarker) {
+            // We keep all resources for snapshots, context options and all events after the marker.
+            visitSha1s(event, sha1s);
+            copyChain = copyChain.then(() => fs.promises.appendFile(traceFileCopy, line + '\n'));
+          }
+        } catch (e) {
+          reject(e);
+          fileStream.close();
+          rl.close();
+        }
+      });
+      rl.on('error', reject);
+      rl.on('close', async () => {
+        await copyChain;
+        resolve();
+      });
+    });
+    return { file: traceFileCopy, sha1s };
   }
 
   async _captureSnapshot(name: 'before' | 'after' | 'action' | 'event', sdkObject: SdkObject, metadata: CallMetadata, element?: ElementHandle) {
@@ -194,10 +259,11 @@ export class Tracing implements InstrumentationListener {
             height: params.height,
             timestamp: monotonicTime()
           };
-          this._appendTraceEvent(event);
-          this._appendEventChain = this._appendEventChain.then(async () => {
+          // Make sure to write the screencast frame before adding a reference to it.
+          this._appendTraceOperation(async () => {
             await fs.promises.writeFile(path.join(this._resourcesDir!, sha1), params.buffer).catch(() => {});
           });
+          this._appendTraceEvent(event);
         }),
     );
   }
@@ -205,31 +271,46 @@ export class Tracing implements InstrumentationListener {
   private _appendTraceEvent(event: any) {
     if (!this._recordingTraceEvents)
       return;
-
-    const visit = (object: any) => {
-      if (Array.isArray(object)) {
-        object.forEach(visit);
-        return;
-      }
-      if (typeof object === 'object') {
-        for (const key in object) {
-          if (key === 'sha1' || key.endsWith('Sha1')) {
-            const sha1 = object[key];
-            if (sha1)
-              this._sha1s.push(sha1);
-          }
-          visit(object[key]);
-        }
-        return;
-      }
-    };
-    visit(event);
-
     // Serialize all writes to the trace file.
-    this._appendEventChain = this._appendEventChain.then(async () => {
+    this._appendTraceOperation(async () => {
+      visitSha1s(event, this._sha1s);
       await fs.promises.appendFile(this._traceFile!, JSON.stringify(event) + '\n');
     });
   }
+
+  private async _appendTraceOperation<T>(cb: () => Promise<T>): Promise<T> {
+    let error: Error | undefined;
+    let result: T | undefined;
+    this._appendEventChain = this._appendEventChain.then(async () => {
+      try {
+        result = await cb();
+      } catch (e) {
+        error = e;
+      }
+    });
+    await this._appendEventChain;
+    if (error)
+      throw error;
+    return result!;
+  }
+}
+
+function visitSha1s(object: any, sha1s: Set<string>) {
+  if (Array.isArray(object)) {
+    object.forEach(o => visitSha1s(o, sha1s));
+    return;
+  }
+  if (typeof object === 'object') {
+    for (const key in object) {
+      if (key === 'sha1' || key.endsWith('Sha1')) {
+        const sha1 = object[key];
+        if (sha1)
+          sha1s.add(sha1);
+      }
+      visitSha1s(object[key], sha1s);
+    }
+    return;
+  }
 }
 
 export function shouldCaptureSnapshot(metadata: CallMetadata): boolean {

tests/tracing.spec.ts

@@ -196,6 +196,27 @@ test('should include interrupted actions', async ({ context, page, server }, tes
   expect(clickEvent.metadata.error.error.message).toBe('Action was interrupted');
 });
 
+test('should reset and export', async ({ context, page, server }, testInfo) => {
+  await context.tracing.start({ screenshots: true, snapshots: true });
+  await page.goto(server.PREFIX + '/frames/frame.html');
+  // @ts-expect-error
+  await context.tracing._reset();
+  await page.setContent('<button>Click</button>');
+  await page.click('"Click"');
+  // @ts-expect-error
+  await context.tracing._export({ path: testInfo.outputPath('trace.zip') });
+  await context.tracing.stop();
+
+  const { events } = await parseTrace(testInfo.outputPath('trace.zip'));
+  expect(events[0].type).toBe('context-options');
+  expect(events.find(e => e.metadata?.apiName === 'page.goto')).toBeFalsy();
+  expect(events.find(e => e.metadata?.apiName === 'page.setContent')).toBeTruthy();
+  expect(events.find(e => e.metadata?.apiName === 'page.click')).toBeTruthy();
+
+  expect(events.some(e => e.type === 'frame-snapshot')).toBeTruthy();
+  expect(events.some(e => e.type === 'resource-snapshot' && e.snapshot.url.endsWith('style.css'))).toBeTruthy();
+});
+
 
 async function parseTrace(file: string): Promise<{ events: any[], resources: Map<string, Buffer> }> {
   const entries = await new Promise<any[]>(f => {