refactor: Update typings for MonkeyPatched (#143)

aws-observability · May 20, 2022 · 3269b13 · 3269b13
1 parent fb6dd84
commit 3269b13
Show file tree

Hide file tree

Showing 7 changed files with 92 additions and 82 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -62,6 +62,7 @@
         "@babel/preset-env": "~7.6.3",
         "@playwright/test": "^1.21.1",
         "@types/jest": "^27.0.2",
+        "@types/shimmer": "^1.0.2",
         "@types/ua-parser-js": "^0.7.35",
         "@types/uuid": "^8.3.0",
         "@webpack-cli/serve": "^1.6.0",

diff --git a/src/plugins/MonkeyPatched.ts b/src/plugins/MonkeyPatched.ts
@@ -1,36 +1,39 @@
 import * as shimmer from 'shimmer';
 
-export type MonkeyPatch = {
-    nodule: object;
-    name: string;
-    // tslint:disable-next-line: ban-types
-    wrapper: Function;
-};
+type Wrapper<W> = () => W;
+export interface MonkeyPatch<
+    Nodule extends object,
+    FieldName extends keyof Nodule
+> {
+    nodule: Nodule;
+    name: FieldName;
+    wrapper: Wrapper<(original: Nodule[FieldName]) => Nodule[FieldName]>;
+}
 
-export abstract class MonkeyPatched {
-    private enabled: boolean;
+export abstract class MonkeyPatched<
+    Nodule extends object,
+    FieldName extends keyof Nodule
+> {
+    public enable = this.patch.bind(this, true);
+    public disable = this.patch.bind(this, false);
 
-    constructor() {
-        this.enabled = false;
-    }
+    protected abstract patches: MonkeyPatch<Nodule, FieldName>[];
 
-    public enable() {
-        if (!this.enabled) {
-            this.enabled = true;
-            for (const patch of this.patches()) {
-                shimmer.wrap(patch.nodule, patch.name, patch.wrapper());
-            }
-        }
-    }
+    private enabled: boolean = false;
 
-    public disable() {
-        if (this.enabled) {
-            this.enabled = false;
-            for (const patch of this.patches()) {
-                shimmer.unwrap(patch.nodule, patch.name);
+    private patch(shouldPatch: boolean = true) {
+        if (this.enabled !== shouldPatch) {
+            this.enabled = shouldPatch;
+            const patchMethod = shouldPatch
+                ? shimmer.wrap.bind(shimmer)
+                : shimmer.unwrap.bind(shimmer);
+            for (const patch of this.patches) {
+                patchMethod(
+                    patch.nodule,
+                    patch.name,
+                    shouldPatch ? patch.wrapper() : undefined
+                );
             }
         }
     }
-
-    protected abstract patches(): MonkeyPatch[];
 }
diff --git a/src/plugins/event-plugins/FetchPlugin.ts b/src/plugins/event-plugins/FetchPlugin.ts
@@ -4,7 +4,7 @@ import {
     Subsegment,
     XRayTraceEvent
 } from '../../events/xray-trace-event';
-import { MonkeyPatch, MonkeyPatched } from '../MonkeyPatched';
+import { MonkeyPatched } from '../MonkeyPatched';
 import {
     PartialHttpPluginConfig,
     defaultConfig,
@@ -47,9 +47,11 @@ export const FETCH_PLUGIN_ID = 'com.amazonaws.rum.fetch';
  * The fetch API is monkey patched using shimmer so all calls to fetch are intercepted. Only calls to URLs which are
  * on the allowlist and are not on the denylist are traced and recorded.
  */
-export class FetchPlugin extends MonkeyPatched implements Plugin {
-    private pluginId: string;
-    private config: HttpPluginConfig;
+export class FetchPlugin
+    extends MonkeyPatched<Window, 'fetch'>
+    implements Plugin {
+    private readonly pluginId: string;
+    private readonly config: HttpPluginConfig;
     private context: PluginContext;
 
     constructor(config?: PartialHttpPluginConfig) {
@@ -67,11 +69,11 @@ export class FetchPlugin extends MonkeyPatched implements Plugin {
         return this.pluginId;
     }
 
-    protected patches(): MonkeyPatch[] {
+    protected get patches() {
         return [
             {
-                nodule: window,
-                name: 'fetch',
+                nodule: window as Window,
+                name: 'fetch' as const,
                 wrapper: this.fetchWrapper
             }
         ];
@@ -277,9 +279,7 @@ export class FetchPlugin extends MonkeyPatched implements Plugin {
             });
     };
 
-    private fetchWrapper = (): ((
-        original: (input: RequestInfo, init?: RequestInit) => Promise<Response>
-    ) => (input: RequestInfo, init?: RequestInit) => Promise<Response>) => {
+    private fetchWrapper = (): ((original: Fetch) => Fetch) => {
         const self = this;
         return (original: Fetch): Fetch => {
             return function (

diff --git a/src/plugins/event-plugins/PageViewPlugin.ts b/src/plugins/event-plugins/PageViewPlugin.ts
@@ -1,5 +1,5 @@
 import { PageIdFormatEnum } from '../../orchestration/Orchestration';
-import { MonkeyPatch, MonkeyPatched } from '../MonkeyPatched';
+import { MonkeyPatched } from '../MonkeyPatched';
 import { Plugin, PluginContext } from '../Plugin';
 
 export const PAGE_EVENT_PLUGIN_ID = 'com.amazonaws.rum.page-view';
@@ -13,8 +13,10 @@ export type Replace = (data: any, title: string, url?: string | null) => void;
  * When a session is initialized, the PageManager records the landing page. When
  * subsequent pages are viewed, this plugin updates the page.
  */
-export class PageViewPlugin extends MonkeyPatched implements Plugin {
-    private pluginId: string;
+export class PageViewPlugin
+    extends MonkeyPatched<History, 'pushState' | 'replaceState'>
+    implements Plugin {
+    private readonly pluginId: string;
     private context: PluginContext;
 
     constructor() {
@@ -33,16 +35,16 @@ export class PageViewPlugin extends MonkeyPatched implements Plugin {
         return this.pluginId;
     }
 
-    protected patches(): MonkeyPatch[] {
+    protected get patches() {
         return [
             {
                 nodule: History.prototype,
-                name: 'pushState',
+                name: 'pushState' as const,
                 wrapper: this.pushState
             },
             {
                 nodule: History.prototype,
-                name: 'replaceState',
+                name: 'replaceState' as const,
                 wrapper: this.replaceState
             }
         ];
@@ -64,7 +66,7 @@ export class PageViewPlugin extends MonkeyPatched implements Plugin {
         };
     };
 
-    private replaceState = (): ((original: Replace) => Replace) => {
+    private replaceState = () => {
         const self = this;
         return (original: Replace): Replace => {
             return function (

diff --git a/src/plugins/event-plugins/XhrPlugin.ts b/src/plugins/event-plugins/XhrPlugin.ts
@@ -18,8 +18,6 @@ import { XhrError } from '../../errors/XhrError';
 import { HTTP_EVENT_TYPE, XRAY_TRACE_EVENT_TYPE } from '../utils/constant';
 import { errorEventToJsErrorEvent } from '../utils/js-error-utils';
 
-type Send = () => void;
-type Open = (method: string, url: string, async: boolean) => void;
 type XhrDetails = {
     method: string;
     url: string;
@@ -93,8 +91,10 @@ export const XHR_PLUGIN_ID = 'com.amazonaws.rum.xhr';
  * - https://xhr.spec.whatwg.org/#event-handlers.
  * - https://xhr.spec.whatwg.org/#events
  */
-export class XhrPlugin extends MonkeyPatched implements Plugin {
-    private pluginId: string;
+export class XhrPlugin
+    extends MonkeyPatched<XMLHttpRequest, 'send' | 'open'>
+    implements Plugin {
+    private readonly pluginId: string;
     private config: HttpPluginConfig;
     private xhrMap: Map<XMLHttpRequest, XhrDetails>;
     private context: PluginContext;
@@ -115,18 +115,18 @@ export class XhrPlugin extends MonkeyPatched implements Plugin {
         return this.pluginId;
     }
 
-    protected patches(): MonkeyPatch[] {
+    protected get patches() {
         return [
             {
                 nodule: XMLHttpRequest.prototype,
-                name: 'send',
+                name: 'send' as const,
                 wrapper: this.sendWrapper
             },
             {
                 nodule: XMLHttpRequest.prototype,
-                name: 'open',
+                name: 'open' as const,
                 wrapper: this.openWrapper
-            }
+            } as MonkeyPatch<XMLHttpRequest, 'open'>
         ];
     }
 
@@ -196,27 +196,17 @@ export class XhrPlugin extends MonkeyPatched implements Plugin {
         const xhr: XMLHttpRequest = e.target as XMLHttpRequest;
         const xhrDetails: XhrDetails = this.xhrMap.get(xhr);
         const errorName: string = 'XMLHttpRequest abort';
-        if (xhrDetails) {
-            const endTime = epochTime();
-            xhrDetails.trace.end_time = endTime;
-            xhrDetails.trace.subsegments[0].end_time = endTime;
-            xhrDetails.trace.subsegments[0].error = true;
-            xhrDetails.trace.subsegments[0].cause = {
-                exceptions: [
-                    {
-                        type: errorName
-                    }
-                ]
-            };
-            this.recordTraceEvent(xhrDetails.trace);
-            this.recordHttpEventWithError(xhrDetails, errorName);
-        }
+        this.handleXhrDetailsOnError(xhrDetails, errorName);
     };
 
     private handleXhrTimeoutEvent = (e: Event) => {
         const xhr: XMLHttpRequest = e.target as XMLHttpRequest;
         const xhrDetails: XhrDetails = this.xhrMap.get(xhr);
         const errorName: string = 'XMLHttpRequest timeout';
+        this.handleXhrDetailsOnError(xhrDetails, errorName);
+    };
+
+    private handleXhrDetailsOnError(xhrDetails: XhrDetails, errorName: string) {
         if (xhrDetails) {
             const endTime = epochTime();
             xhrDetails.trace.end_time = endTime;
@@ -232,7 +222,7 @@ export class XhrPlugin extends MonkeyPatched implements Plugin {
             this.recordTraceEvent(xhrDetails.trace);
             this.recordHttpEventWithError(xhrDetails, errorName);
         }
-    };
+    }
 
     private statusOk(status: number) {
         return status >= 200 && status < 300;
@@ -295,9 +285,9 @@ export class XhrPlugin extends MonkeyPatched implements Plugin {
         );
     };
 
-    private sendWrapper = (): ((original: Send) => Send) => {
+    private sendWrapper = () => {
         const self = this;
-        return (original: Send): Send => {
+        return (original) => {
             return function (this: XMLHttpRequest): void {
                 const xhrDetails: XhrDetails = self.xhrMap.get(this);
                 if (xhrDetails) {
@@ -330,9 +320,9 @@ export class XhrPlugin extends MonkeyPatched implements Plugin {
         };
     };
 
-    private openWrapper = (): ((original: Open) => Open) => {
+    private openWrapper = () => {
         const self = this;
-        return (original: Open): Open => {
+        return (original) => {
             return function (
                 this: XMLHttpRequest,
                 method: string,

diff --git a/src/sessions/VirtualPageLoadTimer.ts b/src/sessions/VirtualPageLoadTimer.ts
@@ -1,20 +1,23 @@
 import { NavigationEvent } from '../events/navigation-event';
 import { PERFORMANCE_NAVIGATION_EVENT_TYPE } from '../plugins/utils/constant';
-import { MonkeyPatch, MonkeyPatched } from '../plugins/MonkeyPatched';
+import { MonkeyPatched } from '../plugins/MonkeyPatched';
 import { Config } from '../orchestration/Orchestration';
 import { RecordEvent } from '../plugins/Plugin';
 import { PageManager, Page } from '../sessions/PageManager';
 
 type Fetch = typeof fetch;
 type Send = () => void;
-
+type Patching = Pick<XMLHttpRequest & Window, 'fetch' | 'send'>;
 /**
  * Maintains the core logic for virtual page load timing functionality.
  * (1) Holds all virtual page load timing related resources
  * (2) Intercepts outgoing XMLHttpRequests and Fetch requests and listens for DOM changes
  * (3) Records virtual page load
  */
-export class VirtualPageLoadTimer extends MonkeyPatched {
+export class VirtualPageLoadTimer extends MonkeyPatched<
+    Patching,
+    'fetch' | 'send'
+> {
     /** Latest interaction time by user on the document */
     public latestInteractionTime: number;
     /** Unique ID of virtual page load periodic checker. */
@@ -98,16 +101,16 @@ export class VirtualPageLoadTimer extends MonkeyPatched {
         this.requestBuffer.clear();
     }
 
-    protected patches(): MonkeyPatch[] {
+    protected get patches() {
         return [
             {
-                nodule: XMLHttpRequest.prototype,
-                name: 'send',
+                nodule: (XMLHttpRequest.prototype as unknown) as Patching,
+                name: 'send' as const,
                 wrapper: this.sendWrapper
             },
             {
-                nodule: window,
-                name: 'fetch',
+                nodule: (window as unknown) as Patching,
+                name: 'fetch' as const,
                 wrapper: this.fetchWrapper
             }
         ];
@@ -173,9 +176,7 @@ export class VirtualPageLoadTimer extends MonkeyPatched {
     /**
      * Increment the fetch counter in PageManager when fetch is beginning
      */
-    private fetchWrapper = (): ((
-        original: (input: RequestInfo, init?: RequestInit) => Promise<Response>
-    ) => (input: RequestInfo, init?: RequestInit) => Promise<Response>) => {
+    private fetchWrapper = (): ((original: Fetch) => Fetch) => {
         const self = this;
         return (original: Fetch): Fetch => {
             return function (