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

packages/core/src/network.js

@@ -1,23 +1,35 @@
 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();
+  // #intercepts = new Map();
   #authentications = new Set();
   #aborted = new Set();
 
@@ -125,7 +137,7 @@
 
     if (!keepPending) {
       this.#pending.delete(requestId);
-      this.#intercepts.delete(requestId);
+      // this.#intercepts.delete(requestId);
     }
   }
 
@@ -153,17 +165,16 @@
   // 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;
@@ -175,16 +186,12 @@
     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);
-      }
     }
+    // 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
@@ -213,8 +220,11 @@
 
   // 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 +234,28 @@
       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 +264,13 @@
   }
 
   // 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/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/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');
+    });
+  });
 });