fix(trace viewer): make LRUCache per-trace (#33245)

packages/trace-viewer/src/sw/lruCache.ts

@@ -0,0 +1,49 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export class LRUCache<K, V> {
+  private _maxSize: number;
+  private _map: Map<K, { value: V, size: number }>;
+  private _size: number;
+
+  constructor(maxSize: number) {
+    this._maxSize = maxSize;
+    this._map = new Map();
+    this._size = 0;
+  }
+
+  getOrCompute(key: K, compute: () => { value: V, size: number }): V {
+    if (this._map.has(key)) {
+      const result = this._map.get(key)!;
+      // reinserting makes this the least recently used entry
+      this._map.delete(key);
+      this._map.set(key, result);
+      return result.value;
+    }
+
+    const result = compute();
+
+    while (this._map.size && this._size + result.size > this._maxSize) {
+      const [firstKey, firstValue] = this._map.entries().next().value;
+      this._size -= firstValue.size;
+      this._map.delete(firstKey);
+    }
+
+    this._map.set(key, result);
+    this._size += result.size;
+    return result.value;
+  }
+}

packages/trace-viewer/src/sw/snapshotRenderer.ts

@@ -17,6 +17,7 @@
 import { escapeHTMLAttribute, escapeHTML } from '@isomorphic/stringUtils';
 import type { FrameSnapshot, NodeNameAttributesChildNodesSnapshot, NodeSnapshot, RenderedFrameSnapshot, ResourceSnapshot, SubtreeReferenceSnapshot } from '@trace/snapshot';
 import type { PageEntry } from '../types/entries';
+import type { LRUCache } from './lruCache';
 
 function findClosest<T>(items: T[], metric: (v: T) => number, target: number) {
   return items.find((item, index) => {
@@ -35,35 +36,8 @@ function isSubtreeReferenceSnapshot(n: NodeSnapshot): n is SubtreeReferenceSnaps
   return Array.isArray(n) && Array.isArray(n[0]);
 }
 
-let cacheSize = 0;
-const cache = new Map<SnapshotRenderer, string>();
-const CACHE_SIZE = 300_000_000; // 300mb
-
-function lruCache(key: SnapshotRenderer, compute: () => string): string {
-  if (cache.has(key)) {
-    const value = cache.get(key)!;
-    // reinserting makes this the least recently used entry
-    cache.delete(key);
-    cache.set(key, value);
-    return value;
-  }
-
-
-  const result = compute();
-
-  while (cache.size && cacheSize + result.length > CACHE_SIZE) {
-    const [firstKey, firstValue] = cache.entries().next().value;
-    cacheSize -= firstValue.length;
-    cache.delete(firstKey);
-  }
-
-  cache.set(key, result);
-  cacheSize += result.length;
-
-  return result;
-}
-
 export class SnapshotRenderer {
+  private _htmlCache: LRUCache<SnapshotRenderer, string>;
   private _snapshots: FrameSnapshot[];
   private _index: number;
   readonly snapshotName: string | undefined;
@@ -72,7 +46,8 @@ export class SnapshotRenderer {
   private _callId: string;
   private _screencastFrames: PageEntry['screencastFrames'];
 
-  constructor(resources: ResourceSnapshot[], snapshots: FrameSnapshot[], screencastFrames: PageEntry['screencastFrames'], index: number) {
+  constructor(htmlCache: LRUCache<SnapshotRenderer, string>, resources: ResourceSnapshot[], snapshots: FrameSnapshot[], screencastFrames: PageEntry['screencastFrames'], index: number) {
+    this._htmlCache = htmlCache;
     this._resources = resources;
     this._snapshots = snapshots;
     this._index = index;
@@ -171,16 +146,15 @@ export class SnapshotRenderer {
     };
 
     const snapshot = this._snapshot;
-    const html = lruCache(this, () => {
+    const html = this._htmlCache.getOrCompute(this, () => {
       visit(snapshot.html, this._index, undefined, undefined);
-
-      const html = result.join('');
-      // Hide the document in order to prevent flickering. We will unhide once script has processed shadow.
       const prefix = snapshot.doctype ? `<!DOCTYPE ${snapshot.doctype}>` : '';
-      return prefix + [
+      const html = prefix + [
+        // Hide the document in order to prevent flickering. We will unhide once script has processed shadow.
         '<style>*,*::before,*::after { visibility: hidden }</style>',
         `<script>${snapshotScript(this._callId, this.snapshotName)}</script>`
-      ].join('') + html;
+      ].join('') + result.join('');
+      return { value: html, size: html.length };
     });
 
     return { html, pageId: snapshot.pageId, frameId: snapshot.frameId, index: this._index };

packages/trace-viewer/src/sw/snapshotStorage.ts

@@ -17,6 +17,7 @@
 import type { FrameSnapshot, ResourceSnapshot } from '@trace/snapshot';
 import { rewriteURLForCustomProtocol, SnapshotRenderer } from './snapshotRenderer';
 import type { PageEntry } from '../types/entries';
+import { LRUCache } from './lruCache';
 
 export class SnapshotStorage {
   private _resources: ResourceSnapshot[] = [];
@@ -24,6 +25,7 @@ export class SnapshotStorage {
     raw: FrameSnapshot[],
     renderers: SnapshotRenderer[]
   }>();
+  private _cache = new LRUCache<SnapshotRenderer, string>(100_000_000);  // 100MB per each trace
 
   addResource(resource: ResourceSnapshot): void {
     resource.request.url = rewriteURLForCustomProtocol(resource.request.url);
@@ -44,7 +46,7 @@ export class SnapshotStorage {
         this._frameSnapshots.set(snapshot.pageId, frameSnapshots);
     }
     frameSnapshots.raw.push(snapshot);
-    const renderer = new SnapshotRenderer(this._resources, frameSnapshots.raw, screencastFrames, frameSnapshots.raw.length - 1);
+    const renderer = new SnapshotRenderer(this._cache, this._resources, frameSnapshots.raw, screencastFrames, frameSnapshots.raw.length - 1);
     frameSnapshots.renderers.push(renderer);
     return renderer;
   }