From 4dcb492ef0c0f25c3f0687638173d20b4b5ded74 Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Mon, 15 Aug 2022 20:54:57 -0700
Subject: [PATCH] fix: make viewer reference all trace urls (#16561)

Single trace viewer page may render several traces, count it as a client for each of the trace files.

Fixes #16429
---
 packages/playwright-core/src/utils/multimap.ts |  7 +++++++
 packages/trace-viewer/src/DEPS.list            |  1 +
 packages/trace-viewer/src/sw.ts                | 11 ++++++-----
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/packages/playwright-core/src/utils/multimap.ts b/packages/playwright-core/src/utils/multimap.ts
index bf9e2870a0..18777d29e5 100644
--- a/packages/playwright-core/src/utils/multimap.ts
+++ b/packages/playwright-core/src/utils/multimap.ts
@@ -46,6 +46,9 @@ export class MultiMap<K, V> {
       this._map.set(key, values.filter(v => value !== v));
   }
 
+  deleteAll(key: K) {
+    this._map.delete(key);
+  }
 
   hasValue(key: K, value: V): boolean {
     const values = this._map.get(key);
@@ -58,6 +61,10 @@ export class MultiMap<K, V> {
     return this._map.size;
   }
 
+  [Symbol.iterator](): Iterator<[K, V[]]> {
+    return this._map[Symbol.iterator]();
+  }
+
   keys(): IterableIterator<K> {
     return this._map.keys();
   }
diff --git a/packages/trace-viewer/src/DEPS.list b/packages/trace-viewer/src/DEPS.list
index 2ebecad755..7f22b2f03d 100644
--- a/packages/trace-viewer/src/DEPS.list
+++ b/packages/trace-viewer/src/DEPS.list
@@ -1,3 +1,4 @@
 [*]
+@playwright-core/utils/multimap.ts
 @web/**
 ui/
diff --git a/packages/trace-viewer/src/sw.ts b/packages/trace-viewer/src/sw.ts
index b030af6b2b..2229ea5fb6 100644
--- a/packages/trace-viewer/src/sw.ts
+++ b/packages/trace-viewer/src/sw.ts
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { MultiMap } from '@playwright-core/utils/multimap';
 import { SnapshotServer } from './snapshotServer';
 import { TraceModel } from './traceModel';
 
@@ -32,11 +33,11 @@ const scopePath = new URL(self.registration.scope).pathname;
 
 const loadedTraces = new Map<string, { traceModel: TraceModel, snapshotServer: SnapshotServer }>();
 
-const clientIdToTraceUrl = new Map<string, string>();
+const clientIdToTraceUrls = new MultiMap<string, string>();
 
 async function loadTrace(trace: string, clientId: string, progress: (done: number, total: number) => void): Promise<TraceModel> {
   const entry = loadedTraces.get(trace);
-  clientIdToTraceUrl.set(clientId, trace);
+  clientIdToTraceUrls.set(clientId, trace);
   if (entry)
     return entry.traceModel;
   const traceModel = new TraceModel();
@@ -123,12 +124,12 @@ async function gc() {
   const clients = await self.clients.matchAll();
   const usedTraces = new Set<string>();
 
-  for (const [clientId, traceUrl] of clientIdToTraceUrl) {
+  for (const [clientId, traceUrls] of clientIdToTraceUrls) {
     // @ts-ignore
     if (!clients.find(c => c.id === clientId))
-      clientIdToTraceUrl.delete(clientId);
+      clientIdToTraceUrls.deleteAll(clientId);
     else
-      usedTraces.add(traceUrl);
+      traceUrls.forEach(url => usedTraces.add(url));
   }
 
   for (const traceUrl of loadedTraces.keys()) {