microsoft · dgozman · Sep 11, 2020 · Sep 11, 2020
diff --git a/src/trace/snapshotter.ts b/src/trace/snapshotter.ts
@@ -18,10 +18,11 @@ import { BrowserContext } from '../server/browserContext';
 import { Page } from '../server/page';
 import * as network from '../server/network';
 import { helper, RegisteredListener } from '../server/helper';
-import { Progress } from '../server/progress';
+import { Progress, runAbortableTask } from '../server/progress';
 import { debugLogger } from '../utils/debugLogger';
 import { Frame } from '../server/frames';
 import * as js from '../server/javascript';
+import * as types from '../server/types';
 import { SnapshotData, takeSnapshotInFrame } from './snapshotterInjected';
 import { assert, calculateSha1, createGuid } from '../utils/utils';
 
@@ -109,32 +110,44 @@ export class Snapshotter {
       this._delegate.onBlob({ sha1, buffer: body });
   }
 
-  async takeSnapshot(progress: Progress, page: Page): Promise<PageSnapshot | null> {
+  async takeSnapshot(page: Page, timeout: number): Promise<PageSnapshot | null> {
     assert(page.context() === this._context);
 
     const frames = page.frames();
-    const promises = frames.map(frame => this._snapshotFrame(progress, frame));
-    const results = await Promise.all(promises);
+    const frameSnapshotPromises = frames.map(async frame => {
+      // TODO: use different timeout depending on the frame depth/origin
+      // to avoid waiting for too long for some useless frame.
+      const frameResult = await runAbortableTask(progress => this._snapshotFrame(progress, frame), timeout).catch(e => null);
+      if (frameResult)
+        return frameResult;
+      const frameSnapshot = {
+        frameId: frame._id,
+        url: frame.url(),
+        html: '<body>Snapshot is not available</body>',
+        resourceOverrides: [],
+      };
+      return { snapshot: frameSnapshot, mapping: new Map<Frame, string>() };
+    });
 
-    const mainFrame = results[0];
-    if (!mainFrame)
+    const viewportSize = await this._getViewportSize(page, timeout);
+    const results = await Promise.all(frameSnapshotPromises);
+
+    if (!viewportSize)
       return null;
+
+    const mainFrame = results[0];
     if (!mainFrame.snapshot.url.startsWith('http'))
       mainFrame.snapshot.url = 'http://playwright.snapshot/';
 
     const mapping = new Map<Frame, string>();
     for (const result of results) {
-      if (!result)
-        continue;
       for (const [key, value] of result.mapping)
         mapping.set(key, value);
     }
 
     const childFrames: FrameSnapshot[] = [];
     for (let i = 1; i < results.length; i++) {
       const result = results[i];
-      if (!result)
-        continue;
       const frame = frames[i];
       if (!mapping.has(frame))
         continue;
@@ -143,75 +156,68 @@ export class Snapshotter {
       childFrames.push(frameSnapshot);
     }
 
-    let viewportSize = page.viewportSize();
-    if (!viewportSize) {
-      try {
-        if (!progress.isRunning())
-          return null;
-
-        const context = await page.mainFrame()._utilityContext();
-        viewportSize = await context.evaluateInternal(() => {
-          return {
-            width: Math.max(document.body.offsetWidth, document.documentElement.offsetWidth),
-            height: Math.max(document.body.offsetHeight, document.documentElement.offsetHeight),
-          };
-        });
-      } catch (e) {
-        return null;
-      }
-    }
-
     return {
       viewportSize,
       frames: [mainFrame.snapshot, ...childFrames],
     };
   }
 
+  private async _getViewportSize(page: Page, timeout: number): Promise<types.Size | null> {
+    return runAbortableTask(async progress => {
+      const viewportSize = page.viewportSize();
+      if (viewportSize)
+        return viewportSize;
+      const context = await page.mainFrame()._utilityContext();
+      return context.evaluateInternal(() => {
+        return {
+          width: Math.max(document.body.offsetWidth, document.documentElement.offsetWidth),
+          height: Math.max(document.body.offsetHeight, document.documentElement.offsetHeight),
+        };
+      });
+    }, timeout).catch(e => null);
+  }
+
   private async _snapshotFrame(progress: Progress, frame: Frame): Promise<FrameSnapshotAndMapping | null> {
-    try {
-      if (!progress.isRunning())
-        return null;
-
-      const context = await frame._utilityContext();
-      const guid = createGuid();
-      const removeNoScript = !frame._page.context()._options.javaScriptEnabled;
-      const result = await js.evaluate(context, false /* returnByValue */, takeSnapshotInFrame, guid, removeNoScript) as js.JSHandle;
-      if (!progress.isRunning())
-        return null;
-
-      const properties = await result.getProperties();
-      const data = await properties.get('data')!.jsonValue() as SnapshotData;
-      const frameElements = await properties.get('frameElements')!.getProperties();
-      result.dispose();
-
-      const snapshot: FrameSnapshot = {
-        frameId: frame._id,
-        url: frame.url(),
-        html: data.html,
-        resourceOverrides: [],
-      };
-      const mapping = new Map<Frame, string>();
-
-      for (const { url, content } of data.resourceOverrides) {
-        const buffer = Buffer.from(content);
-        const sha1 = calculateSha1(buffer);
-        this._delegate.onBlob({ sha1, buffer });
-        snapshot.resourceOverrides.push({ url, sha1 });
-      }
-
-      for (let i = 0; i < data.frameUrls.length; i++) {
-        const element = frameElements.get(String(i))!.asElement();
-        if (!element)
-          continue;
-        const frame = await element.contentFrame().catch(e => null);
-        if (frame)
-          mapping.set(frame, data.frameUrls[i]);
-      }
-
-      return { snapshot, mapping };
-    } catch (e) {
+    if (!progress.isRunning())
       return null;
+
+    const context = await frame._utilityContext();
+    const guid = createGuid();
+    const removeNoScript = !frame._page.context()._options.javaScriptEnabled;
+    const result = await js.evaluate(context, false /* returnByValue */, takeSnapshotInFrame, guid, removeNoScript) as js.JSHandle;
+    if (!progress.isRunning())
+      return null;
+
+    const properties = await result.getProperties();
+    const data = await properties.get('data')!.jsonValue() as SnapshotData;
+    const frameElements = await properties.get('frameElements')!.getProperties();
+    result.dispose();
+
+    const snapshot: FrameSnapshot = {
+      frameId: frame._id,
+      url: frame.url(),
+      html: data.html,
+      resourceOverrides: [],
+    };
+    const mapping = new Map<Frame, string>();
+
+    for (const { url, content } of data.resourceOverrides) {
+      const buffer = Buffer.from(content);
+      const sha1 = calculateSha1(buffer);
+      this._delegate.onBlob({ sha1, buffer });
+      snapshot.resourceOverrides.push({ url, sha1 });
     }
+
+    for (let i = 0; i < data.frameUrls.length; i++) {
+      const element = frameElements.get(String(i))!.asElement();
+      if (!element)
+        continue;
+      const frame = await element.contentFrame().catch(e => null);
+      if (frame)
+        mapping.set(frame, data.frameUrls[i]);
+    }
+
+    return { snapshot, mapping };
   }
 }
 

diff --git a/src/trace/tracer.ts b/src/trace/tracer.ts
@@ -23,11 +23,11 @@ import * as fs from 'fs';
 import { calculateSha1, createGuid, mkdirIfNeeded, monotonicTime } from '../utils/utils';
 import { ActionResult, InstrumentingAgent, instrumentingAgents, ActionMetadata } from '../server/instrumentation';
 import { Page } from '../server/page';
-import { Progress, runAbortableTask } from '../server/progress';
 import { Snapshotter } from './snapshotter';
 import * as types from '../server/types';
 import type { ElementHandle } from '../server/dom';
 import { helper, RegisteredListener } from '../server/helper';
+import { DEFAULT_TIMEOUT } from '../utils/timeoutSettings';
 
 const fsWriteFileAsync = util.promisify(fs.writeFile.bind(fs));
 const fsAppendFileAsync = util.promisify(fs.appendFile.bind(fs));
@@ -59,20 +59,28 @@ export class Tracer implements InstrumentingAgent {
   }
 
   async onContextDestroyed(context: BrowserContext): Promise<void> {
-    const contextTracer = this._contextTracers.get(context);
-    if (contextTracer) {
-      await contextTracer.dispose();
-      this._contextTracers.delete(context);
+    try {
+      const contextTracer = this._contextTracers.get(context);
+      if (contextTracer) {
+        await contextTracer.dispose();
+        this._contextTracers.delete(context);
+      }
+    } catch (e) {
+      // Do not throw from instrumentation.
     }
   }
 
   async onAfterAction(result: ActionResult, metadata?: ActionMetadata): Promise<void> {
-    if (!metadata)
-      return;
-    const contextTracer = this._contextTracers.get(metadata.page.context());
-    if (!contextTracer)
-      return;
-    await contextTracer.recordAction(result, metadata);
+    try {
+      if (!metadata)
+        return;
+      const contextTracer = this._contextTracers.get(metadata.page.context());
+      if (!contextTracer)
+        return;
+      await contextTracer.recordAction(result, metadata);
+    } catch (e) {
+      // Do not throw from instrumentation.
+    }
   }
 }
 
@@ -124,33 +132,21 @@ class ContextTracer implements SnapshotterDelegate {
   }
 
   async captureSnapshot(page: Page, options: types.TimeoutOptions & { label?: string } = {}): Promise<void> {
-    await runAbortableTask(async progress => {
-      const label = options.label || 'snapshot';
-      const snapshot = await this._takeSnapshot(progress, page);
-      if (!snapshot)
-        return;
-      const event: ActionTraceEvent = {
-        type: 'action',
-        contextId: this._contextId,
-        action: 'snapshot',
-        label,
-        snapshot,
-      };
-      this._appendTraceEvent(event);
-    }, page._timeoutSettings.timeout(options));
+    const snapshot = await this._takeSnapshot(page, options.timeout);
+    if (!snapshot)
+      return;
+    const event: ActionTraceEvent = {
+      type: 'action',
+      contextId: this._contextId,
+      action: 'snapshot',
+      label: options.label || 'snapshot',
+      snapshot,
+    };
+    this._appendTraceEvent(event);
   }
 
   async recordAction(result: ActionResult, metadata: ActionMetadata) {
-    let snapshot: { sha1: string, duration: number } | undefined;
-    try {
-      // Use 20% of the default timeout.
-      // Never use zero timeout to avoid stalling because of snapshot.
-      const timeout = (metadata.page._timeoutSettings.timeout({}) / 5) || 6000;
-      snapshot = await runAbortableTask(progress => this._takeSnapshot(progress, metadata.page), timeout);
-    } catch (e) {
-      snapshot = undefined;
-    }
-
+    const snapshot = await this._takeSnapshot(metadata.page);
     const event: ActionTraceEvent = {
       type: 'action',
       contextId: this._contextId,
@@ -196,9 +192,14 @@ class ContextTracer implements SnapshotterDelegate {
     return typeof target === 'string' ? target : await target._previewPromise;
   }
 
-  private async _takeSnapshot(progress: Progress, page: Page): Promise<{ sha1: string, duration: number } | undefined> {
+  private async _takeSnapshot(page: Page, timeout: number = 0): Promise<{ sha1: string, duration: number } | undefined> {
+    if (!timeout) {
+      // Never use zero timeout to avoid stalling because of snapshot.
+      // Use 20% of the default timeout.
+      timeout = (page._timeoutSettings.timeout({}) || DEFAULT_TIMEOUT) / 5;
+    }
     const startTime = monotonicTime();
-    const snapshot = await this._snapshotter.takeSnapshot(progress, page);
+    const snapshot = await this._snapshotter.takeSnapshot(page, timeout);
     if (!snapshot)
       return;
     const buffer = Buffer.from(JSON.stringify(snapshot));

diff --git a/src/utils/timeoutSettings.ts b/src/utils/timeoutSettings.ts
@@ -17,7 +17,8 @@
 
 import { isDebugMode } from './utils';
 
-const DEFAULT_TIMEOUT = isDebugMode() ? 0 : 30000;
+export const DEFAULT_TIMEOUT = 30000;
+const TIMEOUT = isDebugMode() ? 0 : DEFAULT_TIMEOUT;
 
 export class TimeoutSettings {
   private _parent: TimeoutSettings | undefined;
@@ -45,7 +46,7 @@ export class TimeoutSettings {
       return this._defaultTimeout;
     if (this._parent)
       return this._parent.navigationTimeout(options);
-    return DEFAULT_TIMEOUT;
+    return TIMEOUT;
   }
 
   timeout(options: { timeout?: number }): number {
@@ -55,12 +56,12 @@ export class TimeoutSettings {
       return this._defaultTimeout;
     if (this._parent)
       return this._parent.timeout(options);
-    return DEFAULT_TIMEOUT;
+    return TIMEOUT;
   }
 
   static timeout(options: { timeout?: number }): number {
     if (typeof options.timeout === 'number')
       return options.timeout;
-    return DEFAULT_TIMEOUT;
+    return TIMEOUT;
   }
 }
diff --git a/test/snapshot.spec.ts b/test/snapshot.spec.ts
@@ -20,5 +20,5 @@ it('should not throw', (test, parameters) => {
   test.skip(!options.TRACING);
 }, async ({page, server, playwright, toImpl}) => {
   await page.goto(server.PREFIX + '/snapshot/snapshot-with-css.html');
-  await (playwright as any).__tracer.captureSnapshot(toImpl(page), { timeout: 5000, label: 'snapshot' });
+  await (playwright as any).__tracer.captureSnapshot(toImpl(page), { label: 'snapshot' });
 });