Merge pull request #390 from ulixee/loader

fix(puppet): better handling of default loader

Author: blakebyrnes

core/lib/FrameEnvironment.ts

@@ -355,7 +355,8 @@ b) Use the UserProfile feature to set cookies for 1 or more domains before they'
     return this.waitForDom(jsPath, options);
   }
 
-  public waitForLoad(status: IPipelineStatus, options?: IWaitForOptions): Promise<void> {
+  public async waitForLoad(status: IPipelineStatus, options?: IWaitForOptions): Promise<void> {
+    await this.isReady;
     return this.navigationsObserver.waitForLoad(status, options);
   }
 

core/test/domRecorder.test.ts

@@ -225,7 +225,7 @@ function sort() {
     const domFrames = domChanges.filter(x => x.tagName === 'IFRAME');
     expect(domFrames).toHaveLength(3);
 
-    await tab.puppetPage.frames[3].waitForLoad();
+    await tab.frameEnvironmentsByPuppetId.get(tab.puppetPage.frames[3].id).isReady;
 
     await state.db.flush();
     const frames = state.db.frames.all();

interfaces/IPuppetFrame.ts

@@ -14,7 +14,7 @@ export interface IPuppetFrame extends ITypedEventEmitter<IPuppetFrameEvents> {
   isDefaultUrl: boolean;
   isAttached(): boolean;
   resolveNodeId(backendNodeId: number): Promise<string>;
-  waitForLoad(eventName?: keyof ILifecycleEvents): Promise<void>;
+  waitForLoad(eventName?: keyof ILifecycleEvents, timeoutMs?: number): Promise<void>;
   waitForLoader(loaderId?: string): Promise<Error | undefined>;
   canEvaluate(isolatedFromWebPageEnvironment: boolean): boolean;
   getFrameElementNodeId(): Promise<string>;

mitm-socket/index.ts

@@ -181,7 +181,10 @@ export default class MitmSocket
 
   private triggerConnectErrorIfNeeded(isExiting = false): void {
     if (this.connectPromise?.isResolved) return;
-    if (!isExiting && !this.connectError) return;
+    if (isExiting && !this.connectError) {
+      this.connectPromise.resolve();
+      return;
+    }
     this.connectPromise?.reject(
       buildConnectError(
         this.connectError ?? `Socket process exited during connect`,

puppet-chrome/lib/BrowserContext.ts

@@ -34,8 +34,7 @@ import TargetInfo = Protocol.Target.TargetInfo;
 
 export class BrowserContext
   extends TypedEventEmitter<IPuppetContextEvents>
-  implements IPuppetContext
-{
+  implements IPuppetContext {
   public logger: IBoundLog;
 
   public workersById = new Map<string, IPuppetWorker>();
@@ -48,6 +47,7 @@ export class BrowserContext
   private pageOptionsByTargetId = new Map<string, IPuppetPageOptions>();
   private readonly createdTargetIds = new Set<string>();
   private creatingTargetPromises: Promise<void>[] = [];
+  private waitForPageAttachedById = new Map<string, Resolvable<Page>>();
   private readonly browser: Browser;
 
   private isClosing = false;
@@ -81,8 +81,8 @@ export class BrowserContext
   public defaultPageInitializationFn: (page: IPuppetPage) => Promise<any> = () => Promise.resolve();
 
   async newPage(options?: IPuppetPageOptions): Promise<Page> {
-    const resolvable = new Resolvable<void>();
-    this.creatingTargetPromises.push(resolvable.promise);
+    const createTargetPromise = new Resolvable<void>();
+    this.creatingTargetPromises.push(createTargetPromise.promise);
 
     const { targetId } = await this.sendWithBrowserDevtoolsSession('Target.createTarget', {
       url: 'about:blank',
@@ -94,29 +94,24 @@ export class BrowserContext
 
     await this.attachToTarget(targetId);
 
-    resolvable.resolve();
-    const idx = this.creatingTargetPromises.indexOf(resolvable.promise);
+    createTargetPromise.resolve();
+    const idx = this.creatingTargetPromises.indexOf(createTargetPromise.promise);
     if (idx >= 0) this.creatingTargetPromises.splice(idx, 1);
 
-    let hasTimedOut = false;
-    const timeout = setTimeout(() => {
-      hasTimedOut = true;
-    }, 10e3).unref();
-
-    // NOTE: flow here interrupts and expects session to attach and call onPageAttached below
-    while (!this.isClosing) {
-      const page = this.pagesById.get(targetId);
-      if (!page) {
-        if (hasTimedOut) throw new Error('Error creating page. Timed out waiting to attach');
-        await new Promise(setImmediate);
-
-        continue;
-      }
-      clearTimeout(timeout);
-      await page.isReady;
-      if (page.isClosed) throw new Error('Page has been closed.');
-      return page;
+    let page = this.pagesById.get(targetId);
+    if (!page) {
+      const pageAttachedPromise = new Resolvable<Page>(
+        60e3,
+        'Error creating page. Timed out waiting to attach',
+      );
+      this.waitForPageAttachedById.set(targetId, pageAttachedPromise);
+      page = await pageAttachedPromise.promise;
+      this.waitForPageAttachedById.delete(targetId);
     }
+
+    await page.isReady;
+    if (page.isClosed) throw new Error('Page has been closed.');
+    return page;
   }
 
   initializePage(page: Page): Promise<any> {
@@ -150,6 +145,7 @@ export class BrowserContext
 
     const page = new Page(devtoolsSession, targetInfo.targetId, this, this.logger, opener);
     this.pagesById.set(page.targetId, page);
+    this.waitForPageAttachedById.get(page.targetId)?.resolve(page);
     await page.isReady;
     this.emit('page', { page });
     return page;
@@ -220,6 +216,9 @@ export class BrowserContext
     if (this.isClosing) return;
     this.isClosing = true;
 
+    for (const waitingPage of this.waitForPageAttachedById.values()) {
+      waitingPage.reject(new CanceledPromiseError('BrowserContext shutting down'));
+    }
     if (this.browser.devtoolsSession.isConnected()) {
       await Promise.all([...this.pagesById.values()].map(x => x.close()));
       await this.sendWithBrowserDevtoolsSession('Target.disposeBrowserContext', {

puppet-chrome/lib/Frame.ts

@@ -73,6 +73,7 @@ export default class Frame extends TypedEventEmitter<IPuppetFrameEvents> impleme
   private loaderIdResolvers = new Map<string, IResolvablePromise<Error | null>>();
   private readonly devtoolsSession: DevtoolsSession;
   private startedLoaderId: string;
+  private defaultLoaderId: string;
   private resolveLoaderTimeout: NodeJS.Timeout;
 
   private get activeLoader(): IResolvablePromise<Error | null> {
@@ -236,11 +237,15 @@ export default class Frame extends TypedEventEmitter<IPuppetFrameEvents> impleme
   public onLoaded(internalFrame: PageFrame): void {
     this.internalFrame = internalFrame;
     this.updateUrl();
-    this.setLoader(internalFrame.loaderId);
-    if (internalFrame.loaderId && this.url) {
-      this.loaderIdResolvers.get(internalFrame.loaderId).resolve();
+    if (!internalFrame.loaderId) return;
+
+    // if we this is the first loader and url is default, this is the first loader
+    if (this.isDefaultUrl && !this.defaultLoaderId && this.loaderIdResolvers.size === 0) {
+      this.defaultLoaderId = internalFrame.loaderId;
     }
-    if (internalFrame.loaderId && internalFrame.unreachableUrl) {
+    this.setLoader(internalFrame.loaderId);
+
+    if (this.url || internalFrame.unreachableUrl) {
       // if this is a loaded frame, just count it as loaded. it shouldn't fail
       this.loaderIdResolvers.get(internalFrame.loaderId).resolve();
     }
@@ -299,9 +304,17 @@ export default class Frame extends TypedEventEmitter<IPuppetFrameEvents> impleme
     }
   }
 
-  public async waitForLoader(loaderId?: string): Promise<Error | null> {
-    const hasLoaderError = await this.loaderIdResolvers.get(loaderId ?? this.activeLoaderId)
-      ?.promise;
+  public async waitForLoader(loaderId?: string, timeoutMs = 60e3): Promise<Error | null> {
+    if (!loaderId) {
+      loaderId = this.activeLoaderId;
+      if (loaderId === this.defaultLoaderId) {
+        // wait for an actual frame to load
+        const frameLoader = await this.waitOn('frame-loader-created', null, timeoutMs);
+        loaderId = frameLoader.loaderId;
+      }
+    }
+
+    const hasLoaderError = await this.loaderIdResolvers.get(loaderId)?.promise;
     if (hasLoaderError) return hasLoaderError;
 
     if (!this.getActiveContextId(false)) {
@@ -348,7 +361,7 @@ export default class Frame extends TypedEventEmitter<IPuppetFrameEvents> impleme
       lifecycle[name] = new Date();
     }
 
-    if (!this.isDefaultUrl) {
+    if (loaderId !== this.defaultLoaderId) {
       this.emit('frame-lifecycle', { frame: this, name, loaderId: pageLoaderId });
     }
   }
@@ -425,9 +438,15 @@ export default class Frame extends TypedEventEmitter<IPuppetFrameEvents> impleme
     };
   }
 
-  public async waitForLoad(event: keyof ILifecycleEvents = 'load'): Promise<void> {
+  public async waitForLoad(
+    event: keyof ILifecycleEvents = 'load',
+    timeoutMs = 30e3,
+  ): Promise<void> {
+    event ??= 'load';
+    timeoutMs ??= 30e3;
+    await this.waitForLoader(null, timeoutMs);
     if (this.lifecycleEvents[event]) return;
-    await this.waitOn('frame-lifecycle', x => x.name === event, 30e3);
+    await this.waitOn('frame-lifecycle', x => x.name === event, timeoutMs);
   }
 
   private setLoader(loaderId: string): void {

puppet-chrome/lib/FramesManager.ts

@@ -92,6 +92,7 @@ export default class FramesManager extends TypedEventEmitter<IPuppetFrameManager
         }
       } catch (error) {
         if (error instanceof CanceledPromiseError) {
+          resolve();
           return;
         }
         reject(error);

puppet/test/TestPage.ts

@@ -10,7 +10,7 @@ export interface ITestPage extends IPuppetPage {
   type(text: string): Promise<void>;
   attachFrame(frameId: string, url: string): Promise<IPuppetFrame>;
   detachFrame(frameId: string): Promise<void>;
-  goto(url: string, waitOnLifecycle?: string): Promise<void>;
+  goto(url: string, waitOnLifecycle?: string, timeoutMs?: number): Promise<void>;
   setContent(content: string): Promise<void>;
   waitForPopup(): Promise<IPuppetPage>;
 }
@@ -81,13 +81,20 @@ export async function goto(
   page: IPuppetPage,
   url: string,
   waitOnLifecycle: 'load' | 'DOMContentLoaded' = 'load',
+  timeoutMs?: number,
 ) {
-  const nav = page.navigate(url);
-  const lifecycle = page.mainFrame.waitOn('frame-lifecycle', ev => ev.name === waitOnLifecycle);
-  await Promise.all([lifecycle, nav, page.mainFrame.waitOn('frame-navigated')]);
+  waitOnLifecycle ??= 'load';
+  timeoutMs ??= 30e3;
+  await Promise.all([
+    page.navigate(url),
+    page.mainFrame.waitOn('frame-navigated', null, timeoutMs),
+  ]);
+  await page.mainFrame.waitForLoad(waitOnLifecycle, timeoutMs);
 }
 
 export async function setContent(page: IPuppetPage, content: string) {
+  // @ts-ignore
+  page.mainFrame.defaultLoaderId = null;
   await page.evaluate(`((content) => {
         window.stop();
         document.open();

puppet/test/load.test.ts

@@ -62,7 +62,7 @@ describe('Load test', () => {
         page = createTestPage(await context.newPage());
         await page.goto(server.url('link.html'));
 
-        const navigate = page.mainFrame.waitOn('frame-navigated');
+        const navigate = page.mainFrame.waitOn('frame-navigated', null, 75e3);
         await page.click('a');
         await navigate;
         expect(page.mainFrame.url).toBe(`${server.crossProcessBaseUrl}/empty.html`);
@@ -71,5 +71,5 @@ describe('Load test', () => {
       }
     });
     await Promise.all(concurrent);
-  }, 60e3);
+  }, 75e3);
 });