feat(snapshots): make cssom overrides efficient (#5218)

- Intercept CSSOM modifications and recalculate overridden css text.
- When css text does not change, use "backwards reference" similar
  to node references.
- Set 'Cache-Control: no-cache' for resources that could be overridden.

src/cli/traceViewer/snapshotServer.ts

@@ -17,7 +17,7 @@
 import * as http from 'http';
 import * as fs from 'fs';
 import * as path from 'path';
-import type { TraceModel, trace } from './traceModel';
+import type { TraceModel, trace, ContextEntry } from './traceModel';
 import { TraceServer } from './traceServer';
 import { NodeSnapshot } from '../../trace/traceTypes';
 
@@ -119,14 +119,23 @@ export class SnapshotServer {
     function serviceWorkerMain(self: any /* ServiceWorkerGlobalScope */) {
       let traceModel: TraceModel;
 
+      type ContextData = {
+        resourcesByUrl: Map<string, trace.NetworkResourceTraceEvent[]>,
+        overridenUrls: Set<string>
+      };
+      const contextToData = new Map<ContextEntry, ContextData>();
+
       function preprocessModel() {
         for (const contextEntry of traceModel.contexts) {
-          contextEntry.resourcesByUrl = new Map();
+          const contextData: ContextData = {
+            resourcesByUrl: new Map(),
+            overridenUrls: new Set(),
+          };
           const appendResource = (event: trace.NetworkResourceTraceEvent) => {
-            let responseEvents = contextEntry.resourcesByUrl.get(event.url);
+            let responseEvents = contextData.resourcesByUrl.get(event.url);
             if (!responseEvents) {
               responseEvents = [];
-              contextEntry.resourcesByUrl.set(event.url, responseEvents);
+              contextData.resourcesByUrl.set(event.url, responseEvents);
             }
             responseEvents.push(event);
           };
@@ -134,9 +143,16 @@ export class SnapshotServer {
             for (const action of pageEntry.actions)
               action.resources.forEach(appendResource);
             pageEntry.resources.forEach(appendResource);
+            for (const snapshots of Object.values(pageEntry.snapshotsByFrameId)) {
+              for (const snapshot of snapshots) {
+                for (const { url } of snapshot.snapshot.resourceOverrides)
+                  contextData.overridenUrls.add(url);
               }
             }
           }
+          contextToData.set(contextEntry, contextData);
+        }
+      }
 
       self.addEventListener('install', function(event: any) {
         event.waitUntil(fetch('/tracemodel').then(async response => {
@@ -249,6 +265,23 @@ export class SnapshotServer {
         return html;
       }
 
+      function findResourceOverride(snapshots: trace.FrameSnapshotTraceEvent[], snapshotIndex: number, url: string): string | undefined {
+        while (true) {
+          const snapshot = snapshots[snapshotIndex].snapshot;
+          const override = snapshot.resourceOverrides.find(o => o.url === url);
+          if (!override)
+            return;
+          if (override.sha1 !== undefined)
+            return override.sha1;
+          if (override.ref === undefined)
+            return;
+          const referenceIndex = snapshotIndex - override.ref!;
+          if (referenceIndex < 0 || referenceIndex >= snapshotIndex)
+            return;
+          snapshotIndex = referenceIndex;
+        }
+      }
+
       async function doFetch(event: any /* FetchEvent */): Promise<Response> {
         try {
           const pathname = new URL(event.request.url).pathname;
@@ -278,6 +311,7 @@ export class SnapshotServer {
         }
         if (!contextEntry || !pageEntry)
           return request.mode === 'navigate' ? respondNotAvailable() : respond404();
+        const contextData = contextToData.get(contextEntry)!;
 
         const frameSnapshots = pageEntry.snapshotsByFrameId[parsed.frameId] || [];
         let snapshotIndex = -1;
@@ -304,7 +338,8 @@ export class SnapshotServer {
         }
 
         let resource: trace.NetworkResourceTraceEvent | null = null;
-        const resourcesWithUrl = contextEntry.resourcesByUrl.get(removeHash(request.url)) || [];
+        const urlWithoutHash = removeHash(request.url);
+        const resourcesWithUrl = contextData.resourcesByUrl.get(urlWithoutHash) || [];
         for (const resourceEvent of resourcesWithUrl) {
           if (resource && resourceEvent.frameId !== parsed.frameId)
             continue;
@@ -314,22 +349,28 @@ export class SnapshotServer {
         }
         if (!resource)
           return respond404();
-        const resourceOverride = snapshotEvent.snapshot.resourceOverrides.find(o => o.url === request.url);
-        const overrideSha1 = resourceOverride ? resourceOverride.sha1 : undefined;
 
-        const response = overrideSha1 ?
+        const overrideSha1 = findResourceOverride(frameSnapshots, snapshotIndex, urlWithoutHash);
-          await fetch(`/resources/${resource.resourceId}/override/${overrideSha1}`) :
+        const fetchUrl = overrideSha1 ?
-          await fetch(`/resources/${resource.resourceId}`);
+          `/resources/${resource.resourceId}/override/${overrideSha1}` :
+          `/resources/${resource.resourceId}`;
+        const fetchedResponse = await fetch(fetchUrl);
+        const headers = new Headers(fetchedResponse.headers);
         // We make a copy of the response, instead of just forwarding,
         // so that response url is not inherited as "/resources/...", but instead
         // as the original request url.
         // Response url turns into resource base uri that is used to resolve
         // relative links, e.g. url(/foo/bar) in style sheets.
-        return new Response(response.body, {
+        if (contextData.overridenUrls.has(urlWithoutHash)) {
-          status: response.status,
+          // No cache, so that we refetch overridden resources.
-          statusText: response.statusText,
+          headers.set('Cache-Control', 'no-cache');
-          headers: response.headers,
+        }
+        const response = new Response(fetchedResponse.body, {
+          status: fetchedResponse.status,
+          statusText: fetchedResponse.statusText,
+          headers,
         });
+        return response;
       }
 
       self.addEventListener('fetch', function(event: any) {

src/cli/traceViewer/traceModel.ts

@@ -29,7 +29,6 @@ export type ContextEntry = {
   created: trace.ContextCreatedTraceEvent;
   destroyed: trace.ContextDestroyedTraceEvent;
   pages: PageEntry[];
-  resourcesByUrl: Map<string, trace.NetworkResourceTraceEvent[]>;
 }
 
 export type VideoEntry = {
@@ -79,7 +78,6 @@ export function readTraceFile(events: trace.TraceEvent[], traceModel: TraceModel
           created: event,
           destroyed: undefined as any,
           pages: [],
-          resourcesByUrl: new Map(),
         });
         break;
       }
@@ -123,19 +121,12 @@ export function readTraceFile(events: trace.TraceEvent[], traceModel: TraceModel
         break;
       }
       case 'resource': {
-        const contextEntry = contextEntries.get(event.contextId)!;
         const pageEntry = pageEntries.get(event.pageId!)!;
         const action = pageEntry.actions[pageEntry.actions.length - 1];
         if (action)
           action.resources.push(event);
         else
           pageEntry.resources.push(event);
-        let responseEvents = contextEntry.resourcesByUrl.get(event.url);
-        if (!responseEvents) {
-          responseEvents = [];
-          contextEntry.resourcesByUrl.set(event.url, responseEvents);
-        }
-        responseEvents.push(event);
         break;
       }
       case 'dialog-opened':

src/cli/traceViewer/traceViewer.ts

@@ -54,7 +54,6 @@ const emptyModel: TraceModel = {
       name: '<empty>',
       filePath: '',
       pages: [],
-      resourcesByUrl: new Map()
     }
   ]
 };

src/trace/snapshotter.ts

@@ -67,10 +67,14 @@ export class Snapshotter {
         resourceOverrides: [],
       };
       for (const { url, content } of data.resourceOverrides) {
+        if (typeof content === 'string') {
           const buffer = Buffer.from(content);
           const sha1 = calculateSha1(buffer);
           this._delegate.onBlob({ sha1, buffer });
           snapshot.resourceOverrides.push({ url, sha1 });
+        } else {
+          snapshot.resourceOverrides.push({ url, ref: content });
+        }
       }
       this._delegate.onFrameSnapshot(source.frame, data.url, snapshot, data.snapshotId);
     });

src/trace/snapshotterInjected.ts

@@ -29,7 +29,11 @@ export type NodeSnapshot =
 export type SnapshotData = {
   doctype?: string,
   html: NodeSnapshot,
-  resourceOverrides: { url: string, content: string }[],
+  resourceOverrides: {
+    url: string,
+    // String is the content. Number is "x snapshots ago", same url.
+    content: string | number,
+  }[],
   viewport: { width: number, height: number },
   url: string,
   snapshotId?: string,
@@ -48,17 +52,19 @@ export function frameSnapshotStreamer() {
   const kScrollTopAttribute = '__playwright_scroll_top_';
   const kScrollLeftAttribute = '__playwright_scroll_left_';
 
-  // Symbols for our own info on Nodes.
+  // Symbols for our own info on Nodes/StyleSheets.
   const kSnapshotFrameId = Symbol('__playwright_snapshot_frameid_');
   const kCachedData = Symbol('__playwright_snapshot_cache_');
   type CachedData = {
     ref?: [number, number], // Previous snapshotNumber and nodeIndex.
     value?: string, // Value for input/textarea elements.
+    cssText?: string, // Text for stylesheets.
+    cssRef?: number, // Previous snapshotNumber for overridden stylesheets.
   };
-  function ensureCachedData(node: Node): CachedData {
+  function ensureCachedData(obj: any): CachedData {
-    if (!(node as any)[kCachedData])
+    if (!obj[kCachedData])
-      (node as any)[kCachedData] = {};
+      obj[kCachedData] = {};
-    return (node as any)[kCachedData];
+    return obj[kCachedData];
   }
 
   const escaped = { '&': '&amp;', '<': '&lt;', '>': '&gt;', '"': '&quot;', '\'': '&#39;' };
@@ -69,32 +75,45 @@ export function frameSnapshotStreamer() {
     return s.replace(/[&<]/ug, char => (escaped as any)[char]);
   }
 
+  function removeHash(url: string) {
+    try {
+      const u = new URL(url);
+      u.hash = '';
+      return u.toString();
+    } catch (e) {
+      return url;
+    }
+  }
+
   class Streamer {
     private _removeNoScript = true;
-    private _needStyleOverrides = false;
     private _timer: NodeJS.Timeout | undefined;
     private _lastSnapshotNumber = 0;
     private _observer: MutationObserver;
+    private _staleStyleSheets = new Set<CSSStyleSheet>();
+    private _allStyleSheetsWithUrlOverride = new Set<CSSStyleSheet>();
+    private _readingStyleSheet = false;  // To avoid invalidating due to our own reads.
 
     constructor() {
-      // TODO: should we also intercept setters like CSSRule.cssText and CSSStyleRule.selectorText?
+      this._interceptNativeMethod(window.CSSStyleSheet.prototype, 'insertRule', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
-      this._interceptNative(window.CSSStyleSheet.prototype, 'insertRule', () => this._needStyleOverrides = true);
+      this._interceptNativeMethod(window.CSSStyleSheet.prototype, 'deleteRule', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
-      this._interceptNative(window.CSSStyleSheet.prototype, 'deleteRule', () => this._needStyleOverrides = true);
+      this._interceptNativeMethod(window.CSSStyleSheet.prototype, 'addRule', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
-      this._interceptNative(window.CSSStyleSheet.prototype, 'addRule', () => this._needStyleOverrides = true);
+      this._interceptNativeMethod(window.CSSStyleSheet.prototype, 'removeRule', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
-      this._interceptNative(window.CSSStyleSheet.prototype, 'removeRule', () => this._needStyleOverrides = true);
+      this._interceptNativeGetter(window.CSSStyleSheet.prototype, 'rules', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
+      this._interceptNativeGetter(window.CSSStyleSheet.prototype, 'cssRules', (sheet: CSSStyleSheet) => this._invalidateStyleSheet(sheet));
 
       this._observer = new MutationObserver(list => this._handleMutations(list));
       const observerConfig = { attributes: true, childList: true, subtree: true, characterData: true };
       this._observer.observe(document, observerConfig);
-      this._interceptNative(window.Element.prototype, 'attachShadow', (node: Node, shadowRoot: ShadowRoot) => {
+      this._interceptNativeMethod(window.Element.prototype, 'attachShadow', (node: Node, shadowRoot: ShadowRoot) => {
-        this._invalidateCache(node);
+        this._invalidateNode(node);
         this._observer.observe(shadowRoot, observerConfig);
       });
 
       this._streamSnapshot();
     }
 
-    private _interceptNative(obj: any, method: string, cb: (thisObj: any, result: any) => void) {
+    private _interceptNativeMethod(obj: any, method: string, cb: (thisObj: any, result: any) => void) {
       const native = obj[method] as Function;
       if (!native)
         return;
@@ -105,7 +124,58 @@ export function frameSnapshotStreamer() {
       };
     }
 
-    private _invalidateCache(node: Node | null) {
+    private _interceptNativeGetter(obj: any, prop: string, cb: (thisObj: any, result: any) => void) {
+      const descriptor = Object.getOwnPropertyDescriptor(obj, prop)!;
+      Object.defineProperty(obj, prop, {
+        ...descriptor,
+        get: function() {
+          const result = descriptor.get!.call(this);
+          cb(this, result);
+          return result;
+        },
+      });
+    }
+
+    private _invalidateStyleSheet(sheet: CSSStyleSheet) {
+      if (this._readingStyleSheet)
+        return;
+      this._staleStyleSheets.add(sheet);
+      if (sheet.href !== null)
+        this._allStyleSheetsWithUrlOverride.add(sheet);
+      if (sheet.ownerNode && sheet.ownerNode.nodeName === 'STYLE')
+        this._invalidateNode(sheet.ownerNode);
+    }
+
+    private _updateStyleElementStyleSheetTextIfNeeded(sheet: CSSStyleSheet): string | undefined {
+      const data = ensureCachedData(sheet);
+      if (this._staleStyleSheets.has(sheet)) {
+        this._staleStyleSheets.delete(sheet);
+        try {
+          data.cssText = this._getSheetText(sheet);
+        } catch (e) {
+          // Sometimes we cannot access cross-origin stylesheets.
+        }
+      }
+      return data.cssText;
+    }
+
+    // Returns either content, ref, or no override.
+    private _updateLinkStyleSheetTextIfNeeded(sheet: CSSStyleSheet, snapshotNumber: number): string | number | undefined {
+      const data = ensureCachedData(sheet);
+      if (this._staleStyleSheets.has(sheet)) {
+        this._staleStyleSheets.delete(sheet);
+        try {
+          data.cssText = this._getSheetText(sheet);
+          data.cssRef = snapshotNumber;
+          return data.cssText;
+        } catch (e) {
+          // Sometimes we cannot access cross-origin stylesheets.
+        }
+      }
+      return data.cssRef === undefined ? undefined : snapshotNumber - data.cssRef;
+    }
+
+    private _invalidateNode(node: Node | null) {
       while (node) {
         ensureCachedData(node).ref = undefined;
         if (node.nodeType === Node.DOCUMENT_FRAGMENT_NODE && (node as ShadowRoot).host)
@@ -117,7 +187,7 @@ export function frameSnapshotStreamer() {
 
     private _handleMutations(list: MutationRecord[]) {
       for (const mutation of list)
-        this._invalidateCache(mutation.target);
+        this._invalidateNode(mutation.target);
     }
 
     markIframe(iframeElement: HTMLIFrameElement | HTMLFrameElement, frameId: string) {
@@ -177,10 +247,15 @@ export function frameSnapshotStreamer() {
     }
 
     private _getSheetText(sheet: CSSStyleSheet): string {
+      this._readingStyleSheet = true;
+      try {
         const rules: string[] = [];
         for (const rule of sheet.cssRules)
           rules.push(rule.cssText);
         return rules.join('\n');
+      } finally {
+        this._readingStyleSheet = false;
+      }
     }
 
     private _captureSnapshot(snapshotId?: string): SnapshotData {
@@ -194,38 +269,9 @@ export function frameSnapshotStreamer() {
         const value = (input as HTMLInputElement | HTMLTextAreaElement).value;
         const data = ensureCachedData(input);
         if (data.value !== value)
-          this._invalidateCache(input);
+          this._invalidateNode(input);
       }
 
-      const styleNodeToStyleSheetText = new Map<Node, string>();
-      const styleSheetUrlToContentOverride = new Map<string, string>();
-
-      const visitStyleSheet = (sheet: CSSStyleSheet) => {
-        // TODO: recalculate these upon changes, and only send them once.
-        if (!this._needStyleOverrides)
-          return;
-
-        try {
-          for (const rule of sheet.cssRules) {
-            if ((rule as CSSImportRule).styleSheet)
-              visitStyleSheet((rule as CSSImportRule).styleSheet);
-          }
-
-          const cssText = this._getSheetText(sheet);
-          if (sheet.ownerNode && sheet.ownerNode.nodeName === 'STYLE') {
-            // Stylesheets with owner STYLE nodes will be rewritten.
-            styleNodeToStyleSheetText.set(sheet.ownerNode, cssText);
-          } else if (sheet.href !== null) {
-            // Other stylesheets will have resource overrides.
-            const base = this._getSheetBase(sheet);
-            const url = this._resolveUrl(base, sheet.href);
-            styleSheetUrlToContentOverride.set(url, cssText);
-          }
-        } catch (e) {
-          // Sometimes we cannot access cross-origin stylesheets.
-        }
-      };
-
       let nodeCounter = 0;
 
       const visit = (node: Node | ShadowRoot): NodeSnapshot | undefined => {
@@ -253,18 +299,19 @@ export function frameSnapshotStreamer() {
           return escapeText(node.nodeValue || '');
 
         if (nodeName === 'STYLE') {
-          const cssText = styleNodeToStyleSheetText.get(node) || node.textContent || '';
+          const sheet = (node as HTMLStyleElement).sheet;
-          return ['style', {}, escapeText(cssText)];
+          let cssText: string | undefined;
+          if (sheet)
+            cssText = this._updateStyleElementStyleSheetTextIfNeeded(sheet);
+          nodeCounter++;  // Compensate for the extra text node in the list.
+          return ['style', {}, escapeText(cssText || node.textContent || '')];
         }
 
         const attrs: { [attr: string]: string } = {};
         const result: NodeSnapshot = [nodeName, attrs];
 
-        if (nodeType === Node.DOCUMENT_FRAGMENT_NODE) {
+        if (nodeType === Node.DOCUMENT_FRAGMENT_NODE)
-          for (const sheet of (node as ShadowRoot).styleSheets)
-            visitStyleSheet(sheet);
           attrs[kShadowAttribute] = 'open';
-        }
 
         if (nodeType === Node.ELEMENT_NODE) {
           const element = node as Element;
@@ -349,14 +396,11 @@ export function frameSnapshotStreamer() {
         return result;
       };
 
-      for (const sheet of doc.styleSheets)
-        visitStyleSheet(sheet);
       const html = doc.documentElement ? visit(doc.documentElement)! : (['html', {}] as NodeSnapshot);
-
+      const result: SnapshotData = {
-      return {
         html,
         doctype: doc.doctype ? doc.doctype.name : undefined,
-        resourceOverrides: Array.from(styleSheetUrlToContentOverride).map(([url, content]) => ({ url, content })),
+        resourceOverrides: [],
         viewport: {
           width: Math.max(doc.body ? doc.body.offsetWidth : 0, doc.documentElement ? doc.documentElement.offsetWidth : 0),
           height: Math.max(doc.body ? doc.body.offsetHeight : 0, doc.documentElement ? doc.documentElement.offsetHeight : 0),
@@ -364,6 +408,19 @@ export function frameSnapshotStreamer() {
         url: location.href,
         snapshotId,
       };
+
+      for (const sheet of this._allStyleSheetsWithUrlOverride) {
+        const content = this._updateLinkStyleSheetTextIfNeeded(sheet, snapshotNumber);
+        if (content === undefined) {
+          // Unable to capture stylsheet contents.
+          continue;
+        }
+        const base = this._getSheetBase(sheet);
+        const url = removeHash(this._resolveUrl(base, sheet.href!));
+        result.resourceOverrides.push({ url, content });
+      }
+
+      return result;
     }
   }
 

src/trace/traceTypes.ts

@@ -152,6 +152,6 @@ export type TraceEvent =
 export type FrameSnapshot = {
   doctype?: string,
   html: NodeSnapshot,
-  resourceOverrides: { url: string, sha1: string }[],
+  resourceOverrides: { url: string, sha1?: string, ref?: number }[],
   viewport: { width: number, height: number },
 };

test/assets/snapshot/snapshot-with-css.html

@@ -7,14 +7,11 @@
   }
 </style>
 <div>hello, world!</div>
-<textarea>Before edit</textarea>
 <div class=root></div>
 <script>
   let shadow;
 
   window.addEventListener('DOMContentLoaded', () => {
-    document.querySelector('textarea').value = 'After edit';
-
     const root = document.querySelector('.root');
     shadow = root.attachShadow({ mode: 'open' });
 
@@ -27,6 +24,11 @@
     imaged.className = 'imaged';
     shadow.appendChild(imaged);
 
+    const textarea = document.createElement('textarea');
+    textarea.textContent = 'Before edit';
+    textarea.style.display = 'block';
+    shadow.appendChild(textarea);
+
     const iframe = document.createElement('iframe');
     iframe.width = '600px';
     iframe.height = '600px';
@@ -35,6 +37,14 @@
   });
 
   window.addEventListener('load', () => {
+    setTimeout(() => {
+      shadow.querySelector('textarea').value = 'After edit';
+
+      for (const rule of document.styleSheets[1].cssRules) {
+        if (rule.cssText.includes('background: cyan'))
+          rule.style.background = 'magenta';
+      }
+
       for (const rule of shadow.styleSheets[0].cssRules) {
         if (rule.styleSheet) {
           for (const rule2 of rule.styleSheet.cssRules) {
@@ -43,5 +53,6 @@
           }
         }
       }
+    }, 500);
   });
 </script>