chore: remove hierarchy of SnapshotStorage (#21853)

packages/playwright-core/src/server/trace/test/inMemorySnapshotter.ts

@@ -18,23 +18,27 @@ import type { BrowserContext } from '../../browserContext';
 import type { Page } from '../../page';
 import type { FrameSnapshot } from '@trace/snapshot';
 import type { SnapshotRenderer } from '../../../../../trace-viewer/src/snapshotRenderer';
-import { BaseSnapshotStorage } from '../../../../../trace-viewer/src/snapshotStorage';
+import { SnapshotStorage } from '../../../../../trace-viewer/src/snapshotStorage';
 import type { SnapshotterBlob, SnapshotterDelegate } from '../recorder/snapshotter';
 import { Snapshotter } from '../recorder/snapshotter';
 import type { ElementHandle } from '../../dom';
 import type { HarTracerDelegate } from '../../har/harTracer';
 import { HarTracer } from '../../har/harTracer';
 import type * as har from '@trace/har';
+import { ManualPromise } from '../../../utils';
 
-export class InMemorySnapshotter extends BaseSnapshotStorage implements SnapshotterDelegate, HarTracerDelegate {
+export class InMemorySnapshotter implements SnapshotterDelegate, HarTracerDelegate {
   private _blobs = new Map<string, Buffer>();
   private _snapshotter: Snapshotter;
   private _harTracer: HarTracer;
+  private _snapshotReadyPromises = new Map<string, ManualPromise<SnapshotRenderer>>();
+  private _storage: SnapshotStorage;
+  private _snapshotCount = 0;
 
   constructor(context: BrowserContext) {
-    super();
     this._snapshotter = new Snapshotter(context, this);
     this._harTracer = new HarTracer(context, null, this, { content: 'attach', includeTraceInfo: true, recordRequestOverrides: false, waitForContentOnStop: false, skipScripts: true });
+    this._storage = new SnapshotStorage();
   }
 
   async initialize(): Promise<void> {
@@ -47,7 +51,6 @@ export class InMemorySnapshotter extends BaseSnapshotStorage implements Snapshot
     await this._harTracer.flush();
     this._harTracer.stop();
     this._harTracer.start();
-    this.clear();
   }
 
   async dispose() {
@@ -57,25 +60,20 @@ export class InMemorySnapshotter extends BaseSnapshotStorage implements Snapshot
   }
 
   async captureSnapshot(page: Page, callId: string, snapshotName: string, element?: ElementHandle): Promise<SnapshotRenderer> {
-    if (this._frameSnapshots.has(snapshotName))
+    if (this._snapshotReadyPromises.has(snapshotName))
       throw new Error('Duplicate snapshot name: ' + snapshotName);
 
     this._snapshotter.captureSnapshot(page, callId, snapshotName, element).catch(() => {});
-    return new Promise<SnapshotRenderer>(fulfill => {
+    const promise = new ManualPromise<SnapshotRenderer>();
-      const disposable = this.onSnapshotEvent((renderer: SnapshotRenderer) => {
+    this._snapshotReadyPromises.set(snapshotName, promise);
-        if (renderer.snapshotName === snapshotName) {
+    return promise;
-          disposable.dispose();
-          fulfill(renderer);
-        }
-      });
-    });
   }
 
   onEntryStarted(entry: har.Entry) {
   }
 
   onEntryFinished(entry: har.Entry) {
-    this.addResource(entry);
+    this._storage.addResource(entry);
   }
 
   onContentBlob(sha1: string, buffer: Buffer) {
@@ -87,14 +85,16 @@ export class InMemorySnapshotter extends BaseSnapshotStorage implements Snapshot
   }
 
   onFrameSnapshot(snapshot: FrameSnapshot): void {
-    this.addFrameSnapshot(snapshot);
+    ++this._snapshotCount;
-  }
+    const renderer = this._storage.addFrameSnapshot(snapshot);
-
+    this._snapshotReadyPromises.get(snapshot.snapshotName || '')?.resolve(renderer);
-  async resourceContent(sha1: string): Promise<Blob | undefined> {
-    throw new Error('Not implemented');
   }
 
   async resourceContentForTest(sha1: string): Promise<Buffer | undefined> {
     return this._blobs.get(sha1);
   }
+
+  snapshotCount() {
+    return this._snapshotCount;
+  }
 }

packages/trace-viewer/src/snapshotRenderer.ts

@@ -20,7 +20,7 @@ export class SnapshotRenderer {
   private _snapshots: FrameSnapshot[];
   private _index: number;
   readonly snapshotName: string | undefined;
-  _resources: ResourceSnapshot[];
+  private _resources: ResourceSnapshot[];
   private _snapshot: FrameSnapshot;
   private _callId: string;
 

packages/trace-viewer/src/snapshotServer.ts

@@ -14,19 +14,21 @@
  * limitations under the License.
  */
 
-import type { SnapshotStorage } from './snapshotStorage';
 import type { URLSearchParams } from 'url';
 import type { SnapshotRenderer } from './snapshotRenderer';
+import type { SnapshotStorage } from './snapshotStorage';
 import type { ResourceSnapshot } from '@trace/snapshot';
 
 type Point = { x: number, y: number };
 
 export class SnapshotServer {
   private _snapshotStorage: SnapshotStorage;
+  private _resourceLoader: (sha1: string) => Promise<Blob | undefined>;
   private _snapshotIds = new Map<string, SnapshotRenderer>();
 
-  constructor(snapshotStorage: SnapshotStorage) {
+  constructor(snapshotStorage: SnapshotStorage, resourceLoader: (sha1: string) => Promise<Blob | undefined>) {
     this._snapshotStorage = snapshotStorage;
+    this._resourceLoader = resourceLoader;
   }
 
   serveSnapshot(pathname: string, searchParams: URLSearchParams, snapshotUrl: string): Response {
@@ -75,7 +77,7 @@ export class SnapshotServer {
       return new Response(null, { status: 404 });
 
     const sha1 = resource.response.content._sha1;
-    const content = sha1 ? await this._snapshotStorage.resourceContent(sha1) || new Blob([]) : new Blob([]);
+    const content = sha1 ? await this._resourceLoader(sha1) || new Blob([]) : new Blob([]);
 
     let contentType = resource.response.content.mimeType;
     const isTextEncoding = /^text\/|^application\/(javascript|json)/.test(contentType);

packages/trace-viewer/src/snapshotStorage.ts

@@ -15,43 +15,28 @@
  */
 
 import type { FrameSnapshot, ResourceSnapshot } from '@trace/snapshot';
-import { EventEmitter } from './events';
 import { rewriteURLForCustomProtocol, SnapshotRenderer } from './snapshotRenderer';
 
-export interface SnapshotStorage {
+export class SnapshotStorage {
-  resources(): ResourceSnapshot[];
+  private _resources: ResourceSnapshot[] = [];
-  resourceContent(sha1: string): Promise<Blob | undefined>;
+  private _frameSnapshots = new Map<string, {
-  snapshotByName(pageOrFrameId: string, snapshotName: string): SnapshotRenderer | undefined;
-  snapshotByIndex(frameId: string, index: number): SnapshotRenderer | undefined;
-}
-
-export abstract class BaseSnapshotStorage  implements SnapshotStorage {
-  protected _resources: ResourceSnapshot[] = [];
-  protected _frameSnapshots = new Map<string, {
     raw: FrameSnapshot[],
-    renderer: SnapshotRenderer[]
+    renderers: SnapshotRenderer[]
   }>();
-  private _didSnapshot = new EventEmitter<SnapshotRenderer>();
-  readonly onSnapshotEvent = this._didSnapshot.event;
-
-  clear() {
-    this._resources = [];
-    this._frameSnapshots.clear();
-  }
 
   addResource(resource: ResourceSnapshot): void {
     resource.request.url = rewriteURLForCustomProtocol(resource.request.url);
     this._resources.push(resource);
   }
 
-  addFrameSnapshot(snapshot: FrameSnapshot): void {
+  addFrameSnapshot(snapshot: FrameSnapshot) {
     for (const override of snapshot.resourceOverrides)
       override.url = rewriteURLForCustomProtocol(override.url);
     let frameSnapshots = this._frameSnapshots.get(snapshot.frameId);
     if (!frameSnapshots) {
       frameSnapshots = {
         raw: [],
-        renderer: [],
+        renderers: [],
       };
       this._frameSnapshots.set(snapshot.frameId, frameSnapshots);
       if (snapshot.isMainFrame)
@@ -59,23 +44,12 @@ export abstract class BaseSnapshotStorage  implements SnapshotStorage {
     }
     frameSnapshots.raw.push(snapshot);
     const renderer = new SnapshotRenderer(this._resources, frameSnapshots.raw, frameSnapshots.raw.length - 1);
-    frameSnapshots.renderer.push(renderer);
+    frameSnapshots.renderers.push(renderer);
-    this._didSnapshot.fire(renderer);
+    return renderer;
-  }
-
-  abstract resourceContent(sha1: string): Promise<Blob | undefined>;
-
-  resources(): ResourceSnapshot[] {
-    return this._resources.slice();
   }
 
   snapshotByName(pageOrFrameId: string, snapshotName: string): SnapshotRenderer | undefined {
     const snapshot = this._frameSnapshots.get(pageOrFrameId);
-    return snapshot?.renderer.find(r => r.snapshotName === snapshotName);
+    return snapshot?.renderers.find(r => r.snapshotName === snapshotName);
-  }
-
-  snapshotByIndex(frameId: string, index: number): SnapshotRenderer | undefined {
-    const snapshot = this._frameSnapshots.get(frameId);
-    return snapshot?.renderer[index];
   }
 }

packages/trace-viewer/src/sw.ts

@@ -49,7 +49,7 @@ async function loadTrace(traceUrl: string, traceFileName: string | null, clientI
     else
       throw new Error(`Could not load trace from ${traceUrl}. Make sure a valid Playwright Trace is accessible over this url.`);
   }
-  const snapshotServer = new SnapshotServer(traceModel.storage());
+  const snapshotServer = new SnapshotServer(traceModel.storage(), sha1 => traceModel.resourceForSha1(sha1));
   loadedTraces.set(traceUrl, { traceModel, snapshotServer });
   return traceModel;
 }

packages/trace-viewer/src/traceModel.ts

@@ -22,14 +22,14 @@ import type zip from '@zip.js/zip.js';
 import zipImport from '@zip.js/zip.js/dist/zip-no-worker-inflate.min.js';
 import type { ContextEntry, PageEntry } from './entries';
 import { createEmptyContext } from './entries';
-import { BaseSnapshotStorage } from './snapshotStorage';
+import { SnapshotStorage } from './snapshotStorage';
 
 const zipjs = zipImport as typeof zip;
 
 export class TraceModel {
   contextEntries: ContextEntry[] = [];
   pageEntries = new Map<string, PageEntry>();
-  private _snapshotStorage: BaseSnapshotStorage | undefined;
+  private _snapshotStorage: SnapshotStorage | undefined;
   private _version: number | undefined;
   private _backend!: TraceModelBackend;
 
@@ -52,7 +52,7 @@ export class TraceModel {
     if (!ordinals.length)
       throw new Error('Cannot find .trace file');
 
-    this._snapshotStorage = new PersistentSnapshotStorage(this._backend);
+    this._snapshotStorage = new SnapshotStorage();
 
     for (const ordinal of ordinals) {
       const contextEntry = createEmptyContext();
@@ -95,7 +95,7 @@ export class TraceModel {
     return this._backend.readBlob('resources/' + sha1);
   }
 
-  storage(): BaseSnapshotStorage {
+  storage(): SnapshotStorage {
     return this._snapshotStorage!;
   }
 
@@ -387,19 +387,6 @@ class FetchTraceModelBackend implements TraceModelBackend {
   }
 }
 
-export class PersistentSnapshotStorage extends BaseSnapshotStorage {
-  private _backend: TraceModelBackend;
-
-  constructor(backend: TraceModelBackend) {
-    super();
-    this._backend = backend;
-  }
-
-  async resourceContent(sha1: string): Promise<Blob | undefined> {
-    return this._backend.readBlob('resources/' + sha1);
-  }
-}
-
 function formatUrl(trace: string) {
   let url = trace.startsWith('http') || trace.startsWith('blob') ? trace : `file?path=${trace}`;
   // Dropbox does not support cors.

tests/library/snapshotter.spec.ts

@@ -57,11 +57,9 @@ it.describe('snapshots', () => {
 
   it('should collect multiple', async ({ page, toImpl, snapshotter }) => {
     await page.setContent('<button>Hello</button>');
-    const snapshots = [];
-    snapshotter.onSnapshotEvent(snapshot => snapshots.push(snapshot));
     await snapshotter.captureSnapshot(toImpl(page), 'call@1', 'snapshot@call@1');
     await snapshotter.captureSnapshot(toImpl(page), 'call@2', 'snapshot@call@2');
-    expect(snapshots.length).toBe(2);
+    expect(snapshotter.snapshotCount()).toBe(2);
   });
 
   it('should respect inline CSSOM change', async ({ page, toImpl, snapshotter }) => {
@@ -88,7 +86,7 @@ it.describe('snapshots', () => {
     const snapshot1 = await snapshotter.captureSnapshot(toImpl(page), 'call@1', 'snapshot@call@1');
     expect(distillSnapshot(snapshot1)).toBe('<DIV id=\"div\" attr1=\"1\" attr2=\"2\"></DIV>');
     await page.evaluate(() => document.getElementById('div').removeAttribute('attr2'));
-    const snapshot2 = await snapshotter.captureSnapshot(toImpl(page), 'snapshot@call@2');
+    const snapshot2 = await snapshotter.captureSnapshot(toImpl(page), 'call@2', 'snapshot@call@2');
     expect(distillSnapshot(snapshot2)).toBe('<DIV id=\"div\" attr1=\"1\"></DIV>');
   });
 
@@ -125,7 +123,7 @@ it.describe('snapshots', () => {
     expect(distillSnapshot(snapshot1)).toBe('<LINK rel=\"stylesheet\" href=\"style.css\"><BUTTON>Hello</BUTTON>');
 
     await page.evaluate(() => { (document.styleSheets[0].cssRules[0] as any).style.color = 'blue'; });
-    const snapshot2 = await snapshotter.captureSnapshot(toImpl(page), 'call@1', 'snapshot@call@1');
+    const snapshot2 = await snapshotter.captureSnapshot(toImpl(page), 'call@2', 'snapshot@call@2');
     const resource = snapshot2.resourceByUrl(`http://localhost:${server.PORT}/style.css`);
     expect((await snapshotter.resourceContentForTest(resource.response.content._sha1)).toString()).toBe('button { color: blue; }');
   });
@@ -171,7 +169,7 @@ it.describe('snapshots', () => {
 
     // Marking iframe hierarchy is racy, do not expect snapshot, wait for it.
     for (let counter = 0; ; ++counter) {
-      const snapshot = await snapshotter.captureSnapshot(toImpl(page), 'snapshot' + counter);
+      const snapshot = await snapshotter.captureSnapshot(toImpl(page), 'call@' + counter, 'snapshot@call@' + counter);
       const text = distillSnapshot(snapshot).replace(/frame@[^"]+["]/, '<id>"');
       if (text === '<IFRAME __playwright_src__=\"/snapshot/<id>\"></IFRAME>')
         break;
@@ -227,7 +225,7 @@ it.describe('snapshots', () => {
     }
     await handle.evaluate(element => element.setAttribute('data', 'two'));
     {
-      const snapshot = await snapshotter.captureSnapshot(toImpl(page), 'snapshot@call@2');
+      const snapshot = await snapshotter.captureSnapshot(toImpl(page), 'call@3', 'snapshot@call@3');
       expect(distillSnapshot(snapshot)).toBe('<BUTTON data="two">Hello</BUTTON>');
     }
   });
@@ -251,11 +249,11 @@ it.describe('snapshots', () => {
       }
     });
 
-    const renderer1 = await snapshotter.captureSnapshot(toImpl(page), 'call1', 'snapshot@call@1');
+    const renderer1 = await snapshotter.captureSnapshot(toImpl(page), 'call@1', 'snapshot@call@1');
     // Expect some adopted style sheets.
     expect(distillSnapshot(renderer1)).toContain('__playwright_style_sheet_');
 
-    const renderer2 = await snapshotter.captureSnapshot(toImpl(page), 'call2', 'snapshot@call@2');
+    const renderer2 = await snapshotter.captureSnapshot(toImpl(page), 'call@2', 'snapshot@call@2');
     const snapshot2 = renderer2.snapshot();
     // Second snapshot should be just a copy of the first one.
     expect(snapshot2.html).toEqual([[1, 13]]);