chore: merge setStorageState and resetStorage (microsoft#36860)

Author: dgozman

packages/injected/src/storageScript.ts

@@ -176,13 +176,35 @@ export class StorageScript {
     }));
   }
 
-  async restore(originState: channels.SetOriginStorage) {
+  async restore(originState: channels.SetOriginStorage | undefined) {
+    // Clean Service Workers.
+    const registrations = this._global.navigator.serviceWorker ? await this._global.navigator.serviceWorker.getRegistrations() : [];
+    await Promise.all(registrations.map(async r => {
+      // Heuristic for service workers that stalled during main script fetch or importScripts:
+      // Waiting for them to finish unregistering takes ages so we do not await.
+      // However, they will unregister immediately after fetch finishes and should not affect next page load.
+      // Unfortunately, loading next page in Chromium still takes 5 seconds waiting for
+      // some operation on this bogus service worker to finish.
+      if (!r.installing && !r.waiting && !r.active)
+        r.unregister().catch(() => {});
+      else
+        await r.unregister().catch(() => {});
+    }));
+
     try {
-      await Promise.all((originState.indexedDB ?? []).map(dbInfo => this._restoreDB(dbInfo)));
+      for (const db of await this._global.indexedDB.databases?.() || []) {
+        // Do not wait for the callback - it is called on timer in Chromium (slow).
+        if (db.name)
+          this._global.indexedDB.deleteDatabase(db.name!);
+      }
+      await Promise.all((originState?.indexedDB ?? []).map(dbInfo => this._restoreDB(dbInfo)));
     } catch (e) {
       throw new Error('Unable to restore IndexedDB: ' + e.message);
     }
-    for (const { name, value } of (originState.localStorage || []))
+
+    this._global.sessionStorage.clear();
+    this._global.localStorage.clear();
+    for (const { name, value } of (originState?.localStorage || []))
       this._global.localStorage.setItem(name, value);
   }
 }

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

@@ -105,8 +105,7 @@ export abstract class Browser extends SdkObject {
       context._clientCertificatesProxy = clientCertificatesProxy;
       if ((options as any).__testHookBeforeSetStorageState)
         await progress.race((options as any).__testHookBeforeSetStorageState());
-      if (options.storageState)
-        await context.setStorageState(progress, options.storageState);
+      await context.setStorageState(progress, options.storageState, 'initial');
       this.emit(Browser.Events.Context, context);
       return context;
     } catch (error) {

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

@@ -85,7 +85,7 @@ export abstract class BrowserContext extends SdkObject {
   readonly fetchRequest: BrowserContextAPIRequestContext;
   private _customCloseHandler?: () => Promise<any>;
   readonly _tempDirs: string[] = [];
-  private _settingStorageState = false;
+  private _creatingStorageStatePage = false;
   bindingsInitScript?: InitScript;
   initScripts: InitScript[] = [];
   private _routesInFlight = new Set<network.Route>();
@@ -209,15 +209,11 @@ export abstract class BrowserContext extends SdkObject {
 
     // Note: we only need to reset properties from the "paramsThatAllowContextReuse" list.
     // All other properties force a new context.
-    await this._resetStorage(progress);
     await this.clock.uninstall(progress);
     await progress.race(this.setUserAgent(this._options.userAgent));
-    await progress.race(this.clearCache());
-    await progress.race(this.doClearCookies());
     await progress.race(this.doUpdateDefaultEmulatedMedia());
     await progress.race(this.doUpdateDefaultViewport());
-    if (this._options.storageState?.cookies)
-      await progress.race(this.addCookies(this._options.storageState?.cookies));
+    await this.setStorageState(progress, this._options.storageState, 'reset');
 
     await page?.resetForReuse(progress);
   }
@@ -612,68 +608,62 @@ export abstract class BrowserContext extends SdkObject {
     return result;
   }
 
-  async _resetStorage(progress: Progress) {
-    const oldOrigins = this._origins;
-    const newOrigins = new Map(this._options.storageState?.origins?.map(p => [p.origin, p]) || []);
-    if (!oldOrigins.size && !newOrigins.size)
-      return;
-    let page = this.pages()[0];
-
-    // Do not mark this page as internal, because we will leave it for later reuse
-    // as a user-visible page.
-    page = page || await this.newPage(progress, false);
-    const interceptor = (route: network.Route) => {
-      route.fulfill({ body: '<html></html>' }).catch(() => {});
-    };
-    await page.addRequestInterceptor(progress, interceptor, 'prepend');
+  isCreatingStorageStatePage(): boolean {
+    return this._creatingStorageStatePage;
+  }
 
+  async setStorageState(progress: Progress, state: channels.BrowserNewContextParams['storageState'], mode: 'initial' | 'reset') {
+    let page: Page | undefined;
+    let interceptor: network.RouteHandler | undefined;
     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)));
+      if (mode === 'reset') {
+        await progress.race(this.clearCache());
+        await progress.race(this.doClearCookies());
       }
 
-      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);
-    }
-  }
+      if (state?.cookies)
+        await progress.race(this.addCookies(state.cookies));
 
-  isSettingStorageState(): boolean {
-    return this._settingStorageState;
-  }
+      const newOrigins = new Map(state?.origins?.map(p => [p.origin, p]) || []);
+      const allOrigins = new Set([...this._origins, ...newOrigins.keys()]);
+      if (allOrigins.size) {
+        if (mode === 'reset')
+          page = this.pages()[0];
+        if (!page) {
+          try {
+            this._creatingStorageStatePage = mode === 'initial';
+            page = await this.newPage(progress, this._creatingStorageStatePage);
+          } finally {
+            this._creatingStorageStatePage = false;
+          }
+        }
 
-  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)  {
-        page = await this.newPage(progress, true);
-        await page.addRequestInterceptor(progress, route => {
+        interceptor = (route: network.Route) => {
           route.fulfill({ body: '<html></html>' }).catch(() => {});
-        }, 'prepend');
-        for (const originState of state.origins) {
+        };
+        await page.addRequestInterceptor(progress, interceptor, 'prepend');
+
+        for (const origin of allOrigins) {
           const frame = page.mainFrame();
-          await frame.gotoImpl(progress, originState.origin, {});
+          await frame.gotoImpl(progress, origin, {});
           const restoreScript = `(() => {
             const module = {};
             ${rawStorageSource.source}
             const script = new (module.exports.StorageScript())(${this._browser.options.name === 'firefox'});
-            return script.restore(${JSON.stringify(originState)});
+            return script.restore(${JSON.stringify(newOrigins.get(origin))});
           })()`;
           await progress.race(frame.evaluateExpression(restoreScript, { world: 'utility' }));
         }
       }
+      this._origins = new Set([...newOrigins.keys()]);
     } catch (error) {
       rewriteErrorMessage(error, `Error setting storage state:\n` + error.message);
       throw error;
     } finally {
-      await page?.close();
-      this._settingStorageState = false;
+      if (mode === 'initial')
+        await page?.close();
+      else if (interceptor)
+        await page?.removeRequestInterceptor(interceptor);
     }
   }
 

packages/playwright-core/src/server/chromium/crPage.ts

@@ -446,7 +446,7 @@ class FrameSession {
   }
 
   async _initialize(hasUIWindow: boolean) {
-    const isSettingStorageState = this._page.browserContext.isSettingStorageState();
+    const isSettingStorageState = this._page.browserContext.isCreatingStorageStatePage();
     if (!isSettingStorageState && hasUIWindow &&
       !this._crPage._browserContext._browser.isClank() &&
       !this._crPage._browserContext._options.noDefaultViewport) {

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

@@ -1627,40 +1627,6 @@ export class Frame extends SdkObject {
     }, { source, arg });
   }
 
-  async resetStorageForCurrentOriginBestEffort(newStorage: channels.SetOriginStorage | undefined) {
-    const context = await this._utilityContext();
-    await context.evaluate(async ({ ls }) => {
-      // Clean DOMStorage.
-      sessionStorage.clear();
-      localStorage.clear();
-
-      // Add new DOM Storage values.
-      for (const entry of ls || [])
-        localStorage[entry.name] = entry.value;
-
-      // Clean Service Workers
-      const registrations = navigator.serviceWorker ? await navigator.serviceWorker.getRegistrations() : [];
-      await Promise.all(registrations.map(async r => {
-        // Heuristic for service workers that stalled during main script fetch or importScripts:
-        // Waiting for them to finish unregistering takes ages so we do not await.
-        // However, they will unregister immediately after fetch finishes and should not affect next page load.
-        // Unfortunately, loading next page in Chromium still takes 5 seconds waiting for
-        // some operation on this bogus service worker to finish.
-        if (!r.installing && !r.waiting && !r.active)
-          r.unregister().catch(() => {});
-        else
-          await r.unregister().catch(() => {});
-      }));
-
-      // Clean IndexedDB
-      for (const db of await indexedDB.databases?.() || []) {
-        // Do not wait for the callback - it is called on timer in Chromium (slow).
-        if (db.name)
-          indexedDB.deleteDatabase(db.name!);
-      }
-    }, { ls: newStorage?.localStorage }).catch(() => {});
-  }
-
   private _asLocator(selector: string) {
     return asLocator(this._page.browserContext._browser.sdkLanguage(), selector);
   }

packages/playwright-core/src/server/webkit/wkPage.ts

@@ -113,7 +113,7 @@ export class WKPage implements PageDelegate {
   }
 
   private async _initializePageProxySession() {
-    if (this._page.browserContext.isSettingStorageState())
+    if (this._page.browserContext.isCreatingStorageStatePage())
       return;
     const promises: Promise<any>[] = [
       this._pageProxySession.send('Dialog.enable'),
@@ -187,7 +187,7 @@ export class WKPage implements PageDelegate {
       promises.push(session.send('Network.setResourceCachingDisabled', { disabled: true }));
       promises.push(session.send('Network.addInterception', { url: '.*', stage: 'request', isRegex: true }));
     }
-    if (this._page.browserContext.isSettingStorageState()) {
+    if (this._page.browserContext.isCreatingStorageStatePage()) {
       await Promise.all(promises);
       return;
     }