♻️ Refactor: added defaultmap + locking mechanism for intercepts queue (

#1441)

* feat: added defaultmap + locking mechanism for intercepts queue

* feat: added checkResponseSent lock

* chore: add comments + rename variables

* refactor: move DefaultMap to core/utils

* feat: track lifecycle for ALL requests, does not matter if nterwork is intercepted or not

* feat: add locking mechanism for Network.loadingFailed and Network.eventSourceMessageReceived

* refactor: lock names in requestLifeCycle

* chore: lint fix

* fix: coverage issue with single if blocks

* test: added defaultMap specs

* chore: remove unused intercepts map

* Capture service worker intercepted requests (#1443)

* feat: added new config discovery.captureServiceWorker

* feat: use captureServiceWorker config to capture original requests

* chore: lint fix

* test: added spec for captureServiceWorker

* chore: rename captureServiceWorker -> captureMockedServiceWorker

packages/core/src/config.js

@@ -167,6 +167,10 @@ export const configSchema = {
       disableCache: {
         type: 'boolean'
       },
+      captureMockedServiceWorker: {
+        type: 'boolean',
+        default: false
+      },
       requestHeaders: {
         type: 'object',
         normalize: false,
@@ -249,6 +253,7 @@ export const snapshotSchema = {
             requestHeaders: { $ref: '/config/discovery#/properties/requestHeaders' },
             authorization: { $ref: '/config/discovery#/properties/authorization' },
             disableCache: { $ref: '/config/discovery#/properties/disableCache' },
+            captureMockedServiceWorker: { $ref: '/config/discovery#/properties/captureMockedServiceWorker' },
             userAgent: { $ref: '/config/discovery#/properties/userAgent' },
             devicePixelRatio: { $ref: '/config/discovery#/properties/devicePixelRatio' }
           }

packages/core/src/discovery.js

@@ -52,6 +52,7 @@ function debugSnapshotOptions(snapshot) {
   debugProp(snapshot, 'discovery.requestHeaders', JSON.stringify);
   debugProp(snapshot, 'discovery.authorization', JSON.stringify);
   debugProp(snapshot, 'discovery.disableCache');
+  debugProp(snapshot, 'discovery.captureMockedServiceWorker');
   debugProp(snapshot, 'discovery.userAgent');
   debugProp(snapshot, 'clientInfo');
   debugProp(snapshot, 'environmentInfo');
@@ -288,6 +289,7 @@ export function createDiscoveryQueue(percy) {
         requestHeaders: snapshot.discovery.requestHeaders,
         authorization: snapshot.discovery.authorization,
         userAgent: snapshot.discovery.userAgent,
+        captureMockedServiceWorker: snapshot.discovery.captureMockedServiceWorker,
         meta: snapshot.meta,
 
         // enable network inteception

packages/core/src/network.js

@@ -1,23 +1,34 @@
 import { request as makeRequest } from '@percy/client/utils';
 import logger from '@percy/logger';
 import mime from 'mime-types';
-import { createResource, hostnameMatches, normalizeURL, waitFor } from './utils.js';
+import { DefaultMap, createResource, hostnameMatches, normalizeURL, waitFor } from './utils.js';
 
 const MAX_RESOURCE_SIZE = 25 * (1024 ** 2); // 25MB
 const ALLOWED_STATUSES = [200, 201, 301, 302, 304, 307, 308];
 const ALLOWED_RESOURCES = ['Document', 'Stylesheet', 'Image', 'Media', 'Font', 'Other'];
 const ABORTED_MESSAGE = 'Request was aborted by browser';
 
+// RequestLifeCycleHandler handles life cycle of a requestId
+// Ideal flow:          requestWillBeSent -> requestPaused -> responseReceived -> loadingFinished / loadingFailed
+// ServiceWorker flow:  requestWillBeSent -> responseReceived -> loadingFinished / loadingFailed
+class RequestLifeCycleHandler {
+  constructor() {
+    this.resolveRequestWillBeSent = null;
+    this.resolveResponseReceived = null;
+    this.requestWillBeSent = new Promise((resolve) => (this.resolveRequestWillBeSent = resolve));
+    this.responseReceived = new Promise((resolve) => (this.resolveResponseReceived = resolve));
+  }
+}
 // The Interceptor class creates common handlers for dealing with intercepting asset requests
 // for a given page using various devtools protocol events and commands.
 export class Network {
   static TIMEOUT = undefined;
 
   log = logger('core:discovery');
 
+  #requestsLifeCycleHandler = new DefaultMap(() => new RequestLifeCycleHandler());
   #pending = new Map();
   #requests = new Map();
-  #intercepts = new Map();
   #authentications = new Set();
   #aborted = new Set();
 
@@ -26,6 +37,7 @@ export class Network {
     this.timeout = options.networkIdleTimeout ?? 100;
     this.authorization = options.authorization;
     this.requestHeaders = options.requestHeaders ?? {};
+    this.captureMockedServiceWorker = options.captureMockedServiceWorker ?? false;
     this.userAgent = options.userAgent ??
       // by default, emulate a non-headless browser
       page.session.browser.version.userAgent.replace('Headless', '');
@@ -43,7 +55,7 @@ export class Network {
 
     let commands = [
       session.send('Network.enable'),
-      session.send('Network.setBypassServiceWorker', { bypass: true }),
+      session.send('Network.setBypassServiceWorker', { bypass: !this.captureMockedServiceWorker }),
       session.send('Network.setCacheDisabled', { cacheDisabled: true }),
       session.send('Network.setUserAgentOverride', { userAgent: this.userAgent }),
       session.send('Network.setExtraHTTPHeaders', { headers: this.requestHeaders })
@@ -125,7 +137,6 @@ export class Network {
 
     if (!keepPending) {
       this.#pending.delete(requestId);
-      this.#intercepts.delete(requestId);
     }
   }
 
@@ -153,44 +164,42 @@ export class Network {
   // aborted. If the request is already pending, handle it; otherwise set it to be intercepted.
   _handleRequestPaused = async (session, event) => {
     let { networkId: requestId, requestId: interceptId, resourceType } = event;
+
+    // wait for request to be sent
+    await this.#requestsLifeCycleHandler.get(requestId).requestWillBeSent;
     let pending = this.#pending.get(requestId);
     this.#pending.delete(requestId);
 
     // guard against redirects with the same requestId
-    if (pending?.request.url === event.request.url &&
-        pending.request.method === event.request.method) {
-      await this._handleRequest(session, { ...pending, resourceType, interceptId });
-    } else {
-      // track the session that intercepted the request
-      this.#intercepts.set(requestId, { ...event, session });
-    }
+    pending?.request.url === event.request.url &&
+    pending.request.method === event.request.method &&
+    await this._handleRequest(session, { ...pending, resourceType, interceptId });
   }
 
   // Called when a request will be sent. If the request has already been intercepted, handle it;
   // otherwise set it to be pending until it is paused.
   _handleRequestWillBeSent = async event => {
-    let { requestId, request } = event;
+    let { requestId, request, type } = event;
 
     // do not handle data urls
     if (request.url.startsWith('data:')) return;
 
     if (this.intercept) {
-      let intercept = this.#intercepts.get(requestId);
       this.#pending.set(requestId, event);
-
-      if (intercept) {
-        // handle the request with the session that intercepted it
-        let { session, requestId: interceptId, resourceType } = intercept;
-        await this._handleRequest(session, { ...event, resourceType, interceptId });
-        this.#intercepts.delete(requestId);
+      if (this.captureMockedServiceWorker) {
+        await this._handleRequest(undefined, { ...event, resourceType: type, interceptId: requestId }, true);
       }
     }
+    // release request
+    // note: we are releasing this, even if intercept is not set for network.js
+    // since, we want to process all-requests in-order doesn't matter if it should be intercepted or not
+    this.#requestsLifeCycleHandler.get(requestId).resolveRequestWillBeSent();
   }
 
   // Called when a pending request is paused. Handles associating redirected requests with
   // responses and calls this.onrequest with request info and callbacks to continue, respond,
   // or abort a request. One of the callbacks is required to be called and only one.
-  _handleRequest = async (session, event) => {
+  _handleRequest = async (session, event, serviceWorker = false) => {
     let { request, requestId, interceptId, resourceType } = event;
     let redirectChain = [];
 
@@ -208,13 +217,18 @@ export class Network {
     request.redirectChain = redirectChain;
     this.#requests.set(requestId, request);
 
-    await sendResponseResource(this, request, session);
+    if (!serviceWorker) {
+      await sendResponseResource(this, request, session);
+    }
   }
 
   // Called when a response has been received for a specific request. Associates the response with
   // the request data and adds a buffer method to fetch the response body when needed.
-  _handleResponseReceived = (session, event) => {
+  _handleResponseReceived = async (session, event) => {
     let { requestId, response } = event;
+    // await on requestWillBeSent
+    // no explicitly wait on requestWillBePaused as we implictly wait on it, since it manipulates the lifeCycle of request using Fetch module
+    await this.#requestsLifeCycleHandler.get(requestId).requestWillBeSent;
     let request = this.#requests.get(requestId);
     /* istanbul ignore if: race condition paranioa */
     if (!request) return;
@@ -224,20 +238,28 @@ export class Network {
       let result = await this.send(session, 'Network.getResponseBody', { requestId });
       return Buffer.from(result.body, result.base64Encoded ? 'base64' : 'utf-8');
     };
+    // release response
+    this.#requestsLifeCycleHandler.get(requestId).resolveResponseReceived();
   }
 
   // Called when a request streams events. These types of requests break asset discovery because
   // they never finish loading, so we untrack them to signal idle after the first event.
-  _handleEventSourceMessageReceived = event => {
-    let request = this.#requests.get(event.requestId);
+  _handleEventSourceMessageReceived = async event => {
+    let { requestId } = event;
+    // wait for request to be sent
+    await this.#requestsLifeCycleHandler.get(requestId).requestWillBeSent;
+    let request = this.#requests.get(requestId);
     /* istanbul ignore else: race condition paranioa */
     if (request) this._forgetRequest(request);
   }
 
   // Called when a request has finished loading which triggers the this.onrequestfinished
   // callback. The request should have an associated response and be finished with any redirects.
   _handleLoadingFinished = async event => {
-    let request = this.#requests.get(event.requestId);
+    let { requestId } = event;
+    // wait for upto 2 seconds or check if response has been sent
+    await this.#requestsLifeCycleHandler.get(requestId).responseReceived;
+    let request = this.#requests.get(requestId);
     /* istanbul ignore if: race condition paranioa */
     if (!request) return;
 
@@ -246,7 +268,13 @@ export class Network {
   }
 
   // Called when a request has failed loading and triggers the this.onrequestfailed callback.
-  _handleLoadingFailed = event => {
+  _handleLoadingFailed = async event => {
+    let { requestId } = event;
+    // wait for request to be sent
+    // note: we are waiting on requestWillBeSent and NOT responseReceived
+    // since, requests can be cancelled in-flight without Network.responseReceived having been triggered
+    // and in any case, order of processing for responseReceived and loadingFailed does not matter, as response capturing is done in loadingFinished
+    await this.#requestsLifeCycleHandler.get(requestId).requestWillBeSent;
     let request = this.#requests.get(event.requestId);
     /* istanbul ignore if: race condition paranioa */
     if (!request) return;

packages/core/src/snapshot.js

@@ -114,6 +114,7 @@ function getSnapshotOptions(options, { config, meta }) {
       requestHeaders: config.discovery.requestHeaders,
       authorization: config.discovery.authorization,
       disableCache: config.discovery.disableCache,
+      captureMockedServiceWorker: config.discovery.captureMockedServiceWorker,
       userAgent: config.discovery.userAgent
     }
   }, options], (path, prev, next) => {

packages/core/src/utils.js

@@ -320,3 +320,25 @@ export function serializeFunction(fn) {
 
   return fnbody;
 }
+
+// DefaultMap, which returns a default value for an uninitialized key
+// Similar to defaultDict in python
+export class DefaultMap extends Map {
+  constructor(getDefaultValue, ...mapConstructorArgs) {
+    super(...mapConstructorArgs);
+
+    if (typeof getDefaultValue !== 'function') {
+      throw new Error('getDefaultValue must be a function');
+    }
+
+    this.getDefaultValue = getDefaultValue;
+  }
+
+  get = (key) => {
+    if (!this.has(key)) {
+      this.set(key, this.getDefaultValue(key));
+    }
+
+    return super.get(key);
+  };
+};

packages/core/test/discovery.test.js

@@ -1970,4 +1970,68 @@ describe('Discovery', () => {
       });
     });
   });
+
+  describe('Service Worker =>', () => {
+    it('captures original request', async () => {
+      server.reply('/sw.js', () => [200, 'text/javascript', dedent`
+      const fetchUpstream = async(request) => {
+        return await fetch(request.clone());
+      }
+      
+      self.addEventListener('fetch', (event) => {
+        const { request } = event
+        event.respondWith(fetchUpstream(request));
+      });
+      
+      self.addEventListener("activate", (event) => {
+        event.waitUntil(clients.claim());
+      });
+      `]);
+
+      server.reply('/app.js', () => [200, 'text/javascript', dedent`
+      const registerServiceWorker = async () => {
+        await navigator.serviceWorker.register('sw.js',{ scope: './', });
+      };
+      
+      await registerServiceWorker();
+      
+      // create and insert image element which will be intercepted and resolved by service worker
+      // adding a sleep of 1s for service worker to get activated
+      await new Promise(r => setTimeout(r, 1000));
+      var img = document.createElement('img');
+      img.id = 'injected-image';
+      img.src = './img.gif';
+      document.getElementById('container').appendChild(img);
+      `]);
+
+      server.reply('/', () => [200, 'text/html', dedent`
+      <!DOCTYPE html><html><head></head><body>
+        <div id="container"></div>
+        <script type="module" src="app.js"></script>
+      </body></html>
+      `]);
+
+      await percy.snapshot({
+        name: 'first service worker snapshot',
+        url: 'http://localhost:8000',
+        waitForSelector: '#injected-image',
+        discovery: {
+          captureMockedServiceWorker: true
+        }
+      });
+
+      await percy.idle();
+
+      let paths = server.requests.map(r => r[0]);
+      expect(paths).toContain('/img.gif');
+      expect(captured).toContain(jasmine.arrayContaining([
+        jasmine.objectContaining({
+          id: sha256hash(pixel),
+          attributes: jasmine.objectContaining({
+            'resource-url': 'http://localhost:8000/img.gif'
+          })
+        })
+      ]));
+    });
+  });
 });

packages/core/test/unit/utils.test.js

@@ -2,7 +2,8 @@ import {
   generatePromise,
   AbortController,
   yieldTo,
-  yieldAll
+  yieldAll,
+  DefaultMap
 } from '../../src/utils.js';
 
 describe('Unit / Utils', () => {
@@ -165,4 +166,21 @@ describe('Unit / Utils', () => {
         { done: true, value: [2, 4, null, 3, 6] });
     });
   });
+
+  describe('DefaultMap', () => {
+    it('should throw an error if getDefaultValue is not a function', () => {
+      expect(() => new DefaultMap('not a function')).toThrow(new Error('getDefaultValue must be a function'));
+    });
+
+    it('should return the default value for a key that has not been set', () => {
+      const map = new DefaultMap((key) => `default value for ${key}`);
+      expect(map.get('testKey')).toEqual('default value for testKey');
+    });
+
+    it('should return the correct value for a key that has been set', () => {
+      const map = new DefaultMap((key) => `default value for ${key}`);
+      map.set('testKey', 'testValue');
+      expect(map.get('testKey')).toEqual('testValue');
+    });
+  });
 });

packages/core/types/index.d.ts

@@ -18,6 +18,7 @@ interface DiscoveryOptions {
   authorization?: AuthCredentials;
   allowedHostnames?: string[];
   disableCache?: boolean;
+  captureMockedServiceWorker?: boolean;
 }
 
 interface DiscoveryLaunchOptions {