fix(puppet): better handling of default loader

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();
     }
@@ -300,8 +305,16 @@ 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;
+    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, 60e3);
+        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 });
     }
   }

puppet/test/TestPage.ts

@@ -88,6 +88,8 @@ export async function goto(
 }
 
 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);
 });