chore(tracing): fix some of the start/stop scenarios (#6337)

Author: pavelfeldman

src/server/browserContext.ts

@@ -29,7 +29,7 @@ import * as types from './types';
 import path from 'path';
 import { CallMetadata, internalCallMetadata, createInstrumentation, SdkObject } from './instrumentation';
 import { Debugger } from './supplements/debugger';
-import { Tracer } from './trace/recorder/tracer';
+import { Tracing } from './trace/recorder/tracing';
 import { HarTracer } from './supplements/har/harTracer';
 import { RecorderSupplement } from './supplements/recorderSupplement';
 import * as consoleApiSource from '../generated/consoleApiSource';
@@ -57,7 +57,7 @@ export abstract class BrowserContext extends SdkObject {
   private _selectors?: Selectors;
   private _origins = new Set<string>();
   private _harTracer: HarTracer | undefined;
-  readonly tracing: Tracer;
+  readonly tracing: Tracing;
 
   constructor(browser: Browser, options: types.BrowserContextOptions, browserContextId: string | undefined) {
     super(browser, 'browser-context');
@@ -70,7 +70,7 @@ export abstract class BrowserContext extends SdkObject {
 
     if (this._options.recordHar)
       this._harTracer = new HarTracer(this, this._options.recordHar);
-    this.tracing = new Tracer(this);
+    this.tracing = new Tracing(this);
   }
 
   _setSelectors(selectors: Selectors) {
@@ -264,7 +264,7 @@ export abstract class BrowserContext extends SdkObject {
       this._closedStatus = 'closing';
 
       await this._harTracer?.flush();
-      await this.tracing.stop();
+      await this.tracing.dispose();
 
       // Cleanup.
       const promises: Promise<void>[] = [];

src/server/frames.ts

@@ -71,6 +71,7 @@ export class FrameManager {
   readonly _consoleMessageTags = new Map<string, ConsoleTagHandler>();
   readonly _signalBarriers = new Set<SignalBarrier>();
   private _webSockets = new Map<string, network.WebSocket>();
+  readonly _responses: network.Response[] = [];
 
   constructor(page: Page) {
     this._page = page;
@@ -198,6 +199,7 @@ export class FrameManager {
     frame._onClearLifecycle();
     const navigationEvent: NavigationEvent = { url, name, newDocument: frame._currentDocument };
     frame.emit(Frame.Events.Navigation, navigationEvent);
+    this._responses.length = 0;
     if (!initial) {
       debugLogger.log('api', `  navigated to "${url}"`);
       this._page.frameNavigatedToNewDocument(frame);
@@ -264,8 +266,10 @@ export class FrameManager {
   }
 
   requestReceivedResponse(response: network.Response) {
-    if (!response.request()._isFavicon)
-      this._page.emit(Page.Events.Response, response);
+    if (response.request()._isFavicon)
+      return;
+    this._responses.push(response);
+    this._page.emit(Page.Events.Response, response);
   }
 
   requestFinished(request: network.Request) {

src/server/snapshot/inMemorySnapshotter.ts

@@ -38,7 +38,7 @@ export class InMemorySnapshotter extends BaseSnapshotStorage implements Snapshot
   }
 
   async initialize(): Promise<string> {
-    await this._snapshotter.initialize();
+    await this._snapshotter.start();
     return await this._server.start();
   }
 
@@ -62,10 +62,6 @@ export class InMemorySnapshotter extends BaseSnapshotStorage implements Snapshot
     });
   }
 
-  async setAutoSnapshotIntervalForTest(interval: number): Promise<void> {
-    await this._snapshotter.setAutoSnapshotInterval(interval);
-  }
-
   onBlob(blob: SnapshotterBlob): void {
     this._blobs.set(blob.sha1, blob.buffer);
   }

src/server/snapshot/snapshotter.ts

@@ -40,24 +40,46 @@ export class Snapshotter {
   private _context: BrowserContext;
   private _delegate: SnapshotterDelegate;
   private _eventListeners: RegisteredListener[] = [];
-  private _interval = 0;
   private _snapshotStreamer: string;
   private _snapshotBinding: string;
+  private _initialized = false;
+  private _started = false;
+  private _fetchedResponses = new Map<network.Response, string>();
 
   constructor(context: BrowserContext, delegate: SnapshotterDelegate) {
     this._context = context;
     this._delegate = delegate;
-    for (const page of context.pages())
-      this._onPage(page);
-    this._eventListeners = [
-      helper.addEventListener(this._context, BrowserContext.Events.Page, this._onPage.bind(this)),
-    ];
     const guid = createGuid();
     this._snapshotStreamer = '__playwright_snapshot_streamer_' + guid;
     this._snapshotBinding = '__playwright_snapshot_binding_' + guid;
   }
 
-  async initialize() {
+  async start() {
+    this._started = true;
+    if (!this._initialized) {
+      this._initialized = true;
+      await this._initialize();
+    }
+    this._runInAllFrames(`window["${this._snapshotStreamer}"].reset()`);
+
+    // Replay resources loaded in all pages.
+    for (const page of this._context.pages()) {
+      for (const response of page._frameManager._responses)
+        this._saveResource(page, response).catch(e => debugLogger.log('error', e));
+    }
+  }
+
+  async stop() {
+    this._started = false;
+  }
+
+  async _initialize() {
+    for (const page of this._context.pages())
+      this._onPage(page);
+    this._eventListeners = [
+      helper.addEventListener(this._context, BrowserContext.Events.Page, this._onPage.bind(this)),
+    ];
+
     await this._context.exposeBinding(this._snapshotBinding, false, (source, data: SnapshotData) => {
       const snapshot: FrameSnapshot = {
         snapshotName: data.snapshotName,
@@ -87,11 +109,15 @@ export class Snapshotter {
     });
     const initScript = `(${frameSnapshotStreamer})("${this._snapshotStreamer}", "${this._snapshotBinding}")`;
     await this._context._doAddInitScript(initScript);
+    this._runInAllFrames(initScript);
+  }
+
+  private _runInAllFrames(expression: string) {
     const frames = [];
     for (const page of this._context.pages())
       frames.push(...page.frames());
     frames.map(frame => {
-      frame._existingMainContext()?.rawEvaluate(initScript).catch(debugExceptionHandler);
+      frame._existingMainContext()?.rawEvaluate(expression).catch(debugExceptionHandler);
     });
   }
 
@@ -112,37 +138,20 @@ export class Snapshotter {
     page.frames().map(frame => snapshotFrame(frame));
   }
 
-  async setAutoSnapshotInterval(interval: number): Promise<void> {
-    this._interval = interval;
-    const frames = [];
-    for (const page of this._context.pages())
-      frames.push(...page.frames());
-    await Promise.all(frames.map(frame => this._setIntervalInFrame(frame, interval)));
-  }
-
   private _onPage(page: Page) {
-    const processNewFrame = (frame: Frame) => {
-      this._annotateFrameHierarchy(frame);
-      this._setIntervalInFrame(frame, this._interval);
-      const initScript = `(${frameSnapshotStreamer})("${this._snapshotStreamer}", "${this._snapshotBinding}")`;
-      frame._existingMainContext()?.rawEvaluate(initScript).catch(debugExceptionHandler);
-    };
+    // Annotate frame hierarchy so that snapshots could include frame ids.
     for (const frame of page.frames())
-      processNewFrame(frame);
-    this._eventListeners.push(helper.addEventListener(page, Page.Events.FrameAttached, processNewFrame));
-
-    // Push streamer interval on navigation.
-    this._eventListeners.push(helper.addEventListener(page, Page.Events.InternalFrameNavigatedToNewDocument, frame => {
-      this._setIntervalInFrame(frame, this._interval);
-    }));
+      this._annotateFrameHierarchy(frame);
+    this._eventListeners.push(helper.addEventListener(page, Page.Events.FrameAttached, frame => this._annotateFrameHierarchy(frame)));
 
-    // Capture resources.
     this._eventListeners.push(helper.addEventListener(page, Page.Events.Response, (response: network.Response) => {
       this._saveResource(page, response).catch(e => debugLogger.log('error', e));
     }));
   }
 
   private async _saveResource(page: Page, response: network.Response) {
+    if (!this._started)
+      return;
     const isRedirect = response.status() >= 300 && response.status() <= 399;
     if (isRedirect)
       return;
@@ -163,9 +172,25 @@ export class Snapshotter {
     const status = response.status();
     const requestBody = original.postDataBuffer();
     const requestSha1 = requestBody ? calculateSha1(requestBody) : '';
+    if (requestBody)
+      this._delegate.onBlob({ sha1: requestSha1, buffer: requestBody });
     const requestHeaders = original.headers();
-    const body = await response.body().catch(e => debugLogger.log('error', e));
-    const responseSha1 = body ? calculateSha1(body) : '';
+
+    // Only fetch response bodies once.
+    let responseSha1 = this._fetchedResponses.get(response);
+    {
+      if (responseSha1 === undefined) {
+        const body = await response.body().catch(e => debugLogger.log('error', e));
+        // Bail out after each async hop.
+        if (!this._started)
+          return;
+        responseSha1 = body ? calculateSha1(body) : '';
+        if (body)
+          this._delegate.onBlob({ sha1: responseSha1, buffer: body });
+        this._fetchedResponses.set(response, responseSha1);
+      }
+    }
+
     const resource: ResourceSnapshot = {
       pageId: page.guid,
       frameId: response.frame().guid,
@@ -181,17 +206,6 @@ export class Snapshotter {
       timestamp: monotonicTime()
     };
     this._delegate.onResourceSnapshot(resource);
-    if (requestBody)
-      this._delegate.onBlob({ sha1: requestSha1, buffer: requestBody });
-    if (body)
-      this._delegate.onBlob({ sha1: responseSha1, buffer: body });
-  }
-
-  private async _setIntervalInFrame(frame: Frame, interval: number) {
-    const context = frame._existingMainContext();
-    await context?.evaluate(({ snapshotStreamer, interval }) => {
-      (window as any)[snapshotStreamer].setSnapshotInterval(interval);
-    }, { snapshotStreamer: this._snapshotStreamer, interval }).catch(debugExceptionHandler);
   }
 
   private async _annotateFrameHierarchy(frame: Frame) {

src/server/snapshot/snapshotterInjected.ts

@@ -51,6 +51,11 @@ export function frameSnapshotStreamer(snapshotStreamer: string, snapshotBinding:
     cssText?: string, // Text for stylesheets.
     cssRef?: number, // Previous snapshotNumber for overridden stylesheets.
   };
+
+  function resetCachedData(obj: any) {
+    delete obj[kCachedData];
+  }
+
   function ensureCachedData(obj: any): CachedData {
     if (!obj[kCachedData])
       obj[kCachedData] = {};
@@ -69,14 +74,11 @@ export function frameSnapshotStreamer(snapshotStreamer: string, snapshotBinding:
 
   class Streamer {
     private _removeNoScript = true;
-    private _timer: NodeJS.Timeout | undefined;
     private _lastSnapshotNumber = 0;
     private _staleStyleSheets = new Set<CSSStyleSheet>();
-    private _allStyleSheetsWithUrlOverride = new Set<CSSStyleSheet>();
     private _readingStyleSheet = false;  // To avoid invalidating due to our own reads.
     private _fakeBase: HTMLBaseElement;
     private _observer: MutationObserver;
-    private _interval = 0;
 
     constructor() {
       this._interceptNativeMethod(window.CSSStyleSheet.prototype, 'insertRule', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
@@ -125,8 +127,6 @@ export function frameSnapshotStreamer(snapshotStreamer: string, snapshotBinding:
       if (this._readingStyleSheet)
         return;
       this._staleStyleSheets.add(sheet);
-      if (sheet.href !== null)
-        this._allStyleSheetsWithUrlOverride.add(sheet);
     }
 
     private _updateStyleElementStyleSheetTextIfNeeded(sheet: CSSStyleSheet): string | undefined {
@@ -162,29 +162,29 @@ export function frameSnapshotStreamer(snapshotStreamer: string, snapshotBinding:
       (iframeElement as any)[kSnapshotFrameId] = frameId;
     }
 
-    captureSnapshot(snapshotName?: string) {
-      this._streamSnapshot(snapshotName);
-    }
+    reset() {
+      this._staleStyleSheets.clear();
 
-    setSnapshotInterval(interval: number) {
-      this._interval = interval;
-      if (interval)
-        this._streamSnapshot();
+      const visitNode = (node: Node | ShadowRoot) => {
+        resetCachedData(node);
+        if (node.nodeType === Node.ELEMENT_NODE) {
+          const element = node as Element;
+          if (element.shadowRoot)
+            visitNode(element.shadowRoot);
+        }
+        for (let child = node.firstChild; child; child = child.nextSibling)
+          visitNode(child);
+      };
+      visitNode(document.documentElement);
     }
 
-    private _streamSnapshot(snapshotName?: string) {
-      if (this._timer) {
-        clearTimeout(this._timer);
-        this._timer = undefined;
-      }
+    captureSnapshot(snapshotName: string) {
       try {
         const snapshot = this._captureSnapshot(snapshotName);
         if (snapshot)
           (window as any)[snapshotBinding](snapshot);
       } catch (e) {
       }
-      if (this._interval)
-        this._timer = setTimeout(() => this._streamSnapshot(), this._interval);
     }
 
     private _sanitizeUrl(url: string): string {
@@ -298,11 +298,11 @@ export function frameSnapshotStreamer(snapshotStreamer: string, snapshotBinding:
         const result: NodeSnapshot = [nodeName, attrs];
 
         const visitChild = (child: Node) => {
-          const snapshotted = visitNode(child);
-          if (snapshotted) {
-            result.push(snapshotted.n);
+          const snapshot = visitNode(child);
+          if (snapshot) {
+            result.push(snapshot.n);
             expectValue(child);
-            equals = equals && snapshotted.equals;
+            equals = equals && snapshot.equals;
           }
         };
 
@@ -432,10 +432,12 @@ export function frameSnapshotStreamer(snapshotStreamer: string, snapshotBinding:
       };
 
       let allOverridesAreRefs = true;
-      for (const sheet of this._allStyleSheetsWithUrlOverride) {
+      for (const sheet of this._staleStyleSheets) {
+        if (sheet.href === null)
+          continue;
         const content = this._updateLinkStyleSheetTextIfNeeded(sheet, snapshotNumber);
         if (content === undefined) {
-          // Unable to capture stylsheet contents.
+          // Unable to capture stylesheet contents.
           continue;
         }
         if (typeof content !== 'number')

src/server/trace/recorder/traceSnapshotter.ts

@@ -43,7 +43,7 @@ export class TraceSnapshotter extends EventEmitter implements SnapshotterDelegat
   }
 
   async start(): Promise<void> {
-    await this._snapshotter.initialize();
+    await this._snapshotter.start();
   }
 
   async dispose() {