more cleanup, carefully try/catch

Author: dgozman

packages/playwright-core/src/server/android/android.ts

@@ -34,7 +34,7 @@ import { removeFolders } from '../utils/fileUtils';
 import { helper } from '../helper';
 import { SdkObject, serverSideCallMetadata } from '../instrumentation';
 import { gracefullyCloseSet } from '../utils/processLauncher';
-import { isAbortError, Progress, ProgressController } from '../progress';
+import { isAbortError, Progress, ProgressController, raceUncancellableOperationWithCleanup } from '../progress';
 import { registry } from '../registry';
 
 import type { BrowserOptions, BrowserProcess } from '../browser';
@@ -272,8 +272,13 @@ export class AndroidDevice extends SdkObject {
     await progress.race(this._backend.runCommand(`shell:echo "${Buffer.from(commandLine).toString('base64')}" | base64 -d > /data/local/tmp/chrome-command-line`));
     await progress.race(this._backend.runCommand(`shell:am start -a android.intent.action.VIEW -d about:blank ${pkg}`));
     const browserContext = await this._connectToBrowser(progress, socketName, options);
-    await progress.race(this._backend.runCommand(`shell:rm /data/local/tmp/chrome-command-line`));
-    return browserContext;
+    try {
+      await progress.race(this._backend.runCommand(`shell:rm /data/local/tmp/chrome-command-line`));
+      return browserContext;
+    } catch (error) {
+      await browserContext.close({ reason: 'Failed to launch' });
+      throw error;
+    }
   }
 
   private _defaultArgs(options: channels.AndroidDeviceLaunchBrowserParams, socketName: string): string[] {
@@ -314,60 +319,52 @@ export class AndroidDevice extends SdkObject {
 
   private async _connectToBrowser(progress: Progress, socketName: string, options: types.BrowserContextOptions = {}): Promise<BrowserContext> {
     const socket = await this._waitForLocalAbstract(progress, socketName);
-    const androidBrowser = new AndroidBrowser(this, socket);
-    progress.cleanupWhenAborted(() => androidBrowser.close());
-    await progress.race(androidBrowser._init());
-    this._browserConnections.add(androidBrowser);
-
-    const artifactsDir = await progress.race(fs.promises.mkdtemp(ARTIFACTS_FOLDER));
-    const cleanupArtifactsDir = async () => {
-      const errors = (await removeFolders([artifactsDir])).filter(Boolean);
-      for (let i = 0; i < (errors || []).length; ++i)
-        debug('pw:android')(`exception while removing ${artifactsDir}: ${errors[i]}`);
-    };
-    progress.cleanupWhenAborted(cleanupArtifactsDir);
-    gracefullyCloseSet.add(cleanupArtifactsDir);
-    socket.on('close', async () => {
-      gracefullyCloseSet.delete(cleanupArtifactsDir);
-      cleanupArtifactsDir().catch(e => debug('pw:android')(`could not cleanup artifacts dir: ${e}`));
-    });
-    const browserOptions: BrowserOptions = {
-      name: 'clank',
-      isChromium: true,
-      slowMo: 0,
-      persistent: { ...options, noDefaultViewport: true },
-      artifactsDir,
-      downloadsPath: artifactsDir,
-      tracesDir: artifactsDir,
-      browserProcess: new ClankBrowserProcess(androidBrowser),
-      proxy: options.proxy,
-      protocolLogger: helper.debugProtocolLogger(),
-      browserLogsCollector: new RecentLogsCollector(),
-      originalLaunchOptions: {},
-    };
-    validateBrowserContextOptions(options, browserOptions);
-
-    const browser = await progress.race(CRBrowser.connect(this.attribution.playwright, androidBrowser, browserOptions));
-    const defaultContext = browser._defaultContext!;
-    await defaultContext._loadDefaultContextAsIs(progress);
-    return defaultContext;
-  }
-
-  private async _open(progress: Progress, command: string): Promise<SocketBackend> {
-    let aborted = false;
     try {
-      return await progress.race(this._backend.open(command).then(socket => {
-        // Make sure to close the opened socket if progress was aborted before that.
-        if (aborted)
-          socket.close();
-        return socket;
-      }));
+      const androidBrowser = new AndroidBrowser(this, socket);
+      await progress.race(androidBrowser._init());
+      this._browserConnections.add(androidBrowser);
+
+      const artifactsDir = await progress.race(fs.promises.mkdtemp(ARTIFACTS_FOLDER));
+      const cleanupArtifactsDir = async () => {
+        const errors = (await removeFolders([artifactsDir])).filter(Boolean);
+        for (let i = 0; i < (errors || []).length; ++i)
+          debug('pw:android')(`exception while removing ${artifactsDir}: ${errors[i]}`);
+      };
+      gracefullyCloseSet.add(cleanupArtifactsDir);
+      socket.on('close', async () => {
+        gracefullyCloseSet.delete(cleanupArtifactsDir);
+        cleanupArtifactsDir().catch(e => debug('pw:android')(`could not cleanup artifacts dir: ${e}`));
+      });
+      const browserOptions: BrowserOptions = {
+        name: 'clank',
+        isChromium: true,
+        slowMo: 0,
+        persistent: { ...options, noDefaultViewport: true },
+        artifactsDir,
+        downloadsPath: artifactsDir,
+        tracesDir: artifactsDir,
+        browserProcess: new ClankBrowserProcess(androidBrowser),
+        proxy: options.proxy,
+        protocolLogger: helper.debugProtocolLogger(),
+        browserLogsCollector: new RecentLogsCollector(),
+        originalLaunchOptions: {},
+      };
+      validateBrowserContextOptions(options, browserOptions);
+
+      const browser = await progress.race(CRBrowser.connect(this.attribution.playwright, androidBrowser, browserOptions));
+      const defaultContext = browser._defaultContext!;
+      await defaultContext._loadDefaultContextAsIs(progress);
+      return defaultContext;
     } catch (error) {
-      aborted = true;
+      socket.close();
       throw error;
     }
   }
 
+  private _open(progress: Progress, command: string): Promise<SocketBackend> {
+    return raceUncancellableOperationWithCleanup(progress, () => this._backend.open(command), socket => socket.close());
+  }
+
   webViews(): channels.AndroidWebView[] {
     return [...this._webViews.values()];
   }

packages/playwright-core/src/server/browser.ts

@@ -92,11 +92,6 @@ export abstract class Browser extends SdkObject {
     return this.options.sdkLanguage || this.attribution.playwright.options.sdkLanguage;
   }
 
-  newContextFromMetadata(metadata: CallMetadata, options: types.BrowserContextOptions): Promise<BrowserContext> {
-    const controller = new ProgressController(metadata, this);
-    return controller.run(progress => this.newContext(progress, options));
-  }
-
   async newContext(progress: Progress, options: types.BrowserContextOptions): Promise<BrowserContext> {
     validateBrowserContextOptions(options, this.options);
     let clientCertificatesProxy: ClientCertificatesProxy | undefined;

packages/playwright-core/src/server/browserContext.ts

@@ -334,10 +334,14 @@ export abstract class BrowserContext extends SdkObject {
     const binding = new PageBinding(name, playwrightBinding, needsHandle);
     binding.forClient = forClient;
     this._pageBindings.set(name, binding);
-    progress.cleanupWhenAborted(() => this._pageBindings.delete(name));
-    await progress.race(this.doAddInitScript(binding.initScript));
-    await progress.race(this.safeNonStallingEvaluateInAllFrames(binding.initScript.source, 'main'));
-    return binding;
+    try {
+      await progress.race(this.doAddInitScript(binding.initScript));
+      await progress.race(this.safeNonStallingEvaluateInAllFrames(binding.initScript.source, 'main'));
+      return binding;
+    } catch (error) {
+      this._pageBindings.delete(name);
+      throw error;
+    }
   }
 
   async removeExposedBindings(bindings: PageBinding[]) {
@@ -370,21 +374,27 @@ export abstract class BrowserContext extends SdkObject {
   async setExtraHTTPHeaders(progress: Progress, headers: types.HeadersArray) {
     const oldHeaders = this._options.extraHTTPHeaders;
     this._options.extraHTTPHeaders = headers;
-    progress.cleanupWhenAborted(async () => {
+    try {
+      await progress.race(this.doUpdateExtraHTTPHeaders());
+    } catch (error) {
       this._options.extraHTTPHeaders = oldHeaders;
-      await this.doUpdateExtraHTTPHeaders();
-    });
-    await progress.race(this.doUpdateExtraHTTPHeaders());
+      // Note: no await, headers will be reset in the background as soon as possible.
+      this.doUpdateExtraHTTPHeaders().catch(() => {});
+      throw error;
+    }
   }
 
   async setOffline(progress: Progress, offline: boolean) {
     const oldOffline = this._options.offline;
     this._options.offline = offline;
-    progress.cleanupWhenAborted(async () => {
+    try {
+      await progress.race(this.doUpdateOffline());
+    } catch (error) {
       this._options.offline = oldOffline;
-      await this.doUpdateOffline();
-    });
-    await progress.race(this.doUpdateOffline());
+      // Note: no await, offline will be reset in the background as soon as possible.
+      this.doUpdateOffline().catch(() => {});
+      throw error;
+    }
   }
 
   async _loadDefaultContextAsIs(progress: Progress): Promise<Page | undefined> {
@@ -442,13 +452,18 @@ export abstract class BrowserContext extends SdkObject {
   async addInitScript(progress: Progress | undefined, source: string) {
     const initScript = new InitScript(source);
     this.initScripts.push(initScript);
-    progress?.cleanupWhenAborted(() => this.removeInitScripts([initScript]));
-    const promise = this.doAddInitScript(initScript);
-    if (progress)
-      await progress.race(promise);
-    else
-      await promise;
-    return initScript;
+    try {
+      const promise = this.doAddInitScript(initScript);
+      if (progress)
+        await progress.race(promise);
+      else
+        await promise;
+      return initScript;
+    } catch (error) {
+      // Note: no await, init script will be removed in the background as soon as possible.
+      this.removeInitScripts([initScript]).catch(() => {});
+      throw error;
+    }
   }
 
   async removeInitScripts(initScripts: InitScript[]) {
@@ -528,9 +543,8 @@ export abstract class BrowserContext extends SdkObject {
   }
 
   async newPage(progress: Progress, isServerSide: boolean): Promise<Page> {
-    let page: Page | undefined;
+    const page = await progress.race(this.doCreateNewPage(isServerSide));
     try {
-      page = await progress.race(this.doCreateNewPage(isServerSide));
       const pageOrError = await progress.race(page.waitForInitializedOrError());
       if (pageOrError instanceof Page) {
         if (pageOrError.isClosed())
@@ -539,7 +553,7 @@ export abstract class BrowserContext extends SdkObject {
       }
       throw pageOrError;
     } catch (error) {
-      await page?.close({ reason: 'Failed to create page' }).catch(() => {});
+      await page.close({ reason: 'Failed to create page' }).catch(() => {});
       throw error;
     }
   }
@@ -580,17 +594,20 @@ export abstract class BrowserContext extends SdkObject {
     // If there are still origins to save, create a blank page to iterate over origins.
     if (originsToSave.size)  {
       const page = await this.newPage(progress, true);
-      await page.addRequestInterceptor(progress, route => {
-        route.fulfill({ body: '<html></html>' }).catch(() => {});
-      }, 'prepend');
-      for (const origin of originsToSave) {
-        const frame = page.mainFrame();
-        await frame.gotoImpl(progress, origin, {});
-        const storage: SerializedStorage = await progress.race(frame.evaluateExpression(collectScript, { world: 'utility' }));
-        if (storage.localStorage.length || storage.indexedDB?.length)
-          result.origins.push({ origin, localStorage: storage.localStorage, indexedDB: storage.indexedDB });
+      try {
+        await page.addRequestInterceptor(progress, route => {
+          route.fulfill({ body: '<html></html>' }).catch(() => {});
+        }, 'prepend');
+        for (const origin of originsToSave) {
+          const frame = page.mainFrame();
+          await frame.gotoImpl(progress, origin, {});
+          const storage: SerializedStorage = await progress.race(frame.evaluateExpression(collectScript, { world: 'utility' }));
+          if (storage.localStorage.length || storage.indexedDB?.length)
+            result.origins.push({ origin, localStorage: storage.localStorage, indexedDB: storage.indexedDB });
+        }
+      } finally {
+        await page.close();
       }
-      await page.close();
     }
     return result;
   }
@@ -608,33 +625,34 @@ export abstract class BrowserContext extends SdkObject {
     const interceptor = (route: network.Route) => {
       route.fulfill({ body: '<html></html>' }).catch(() => {});
     };
-
-    progress.cleanupWhenAborted(() => page.removeRequestInterceptor(interceptor));
     await page.addRequestInterceptor(progress, interceptor, 'prepend');
 
-    for (const origin of new Set([...oldOrigins, ...newOrigins.keys()])) {
-      const frame = page.mainFrame();
-      await frame.gotoImpl(progress, origin, {});
-      await progress.race(frame.resetStorageForCurrentOriginBestEffort(newOrigins.get(origin)));
-    }
-
-    await page.removeRequestInterceptor(interceptor);
+    try {
+      for (const origin of new Set([...oldOrigins, ...newOrigins.keys()])) {
+        const frame = page.mainFrame();
+        await frame.gotoImpl(progress, origin, {});
+        await progress.race(frame.resetStorageForCurrentOriginBestEffort(newOrigins.get(origin)));
+      }
 
-    this._origins = new Set([...newOrigins.keys()]);
-    // It is safe to not restore the URL to about:blank since we are doing it in Page::resetForReuse.
+      this._origins = new Set([...newOrigins.keys()]);
+      // It is safe to not restore the URL to about:blank since we are doing it in Page::resetForReuse.
+    } finally {
+      await page.removeRequestInterceptor(interceptor);
+    }
   }
 
   isSettingStorageState(): boolean {
     return this._settingStorageState;
   }
 
   async setStorageState(progress: Progress, state: NonNullable<channels.BrowserNewContextParams['storageState']>) {
+    let page: Page | undefined;
     this._settingStorageState = true;
     try {
       if (state.cookies)
         await progress.race(this.addCookies(state.cookies));
       if (state.origins && state.origins.length)  {
-        const page = await this.newPage(progress, true);
+        page = await this.newPage(progress, true);
         await page.addRequestInterceptor(progress, route => {
           route.fulfill({ body: '<html></html>' }).catch(() => {});
         }, 'prepend');
@@ -649,12 +667,12 @@ export abstract class BrowserContext extends SdkObject {
           })()`;
           await progress.race(frame.evaluateExpression(restoreScript, { world: 'utility' }));
         }
-        await page.close();
       }
     } catch (error) {
       rewriteErrorMessage(error, `Error setting storage state:\n` + error.message);
       throw error;
     } finally {
+      await page?.close();
       this._settingStorageState = false;
     }
   }