fix(core): cleanup event listener memory

Author: blakebyrnes

client/package.json

@@ -11,7 +11,7 @@
     "@secret-agent/interfaces": "1.6.2",
     "@secret-agent/plugin-utils": "1.6.2",
     "@secret-agent/replay": "1.6.2",
-    "awaited-dom": "1.3.0",
+    "awaited-dom": "1.3.1",
     "uuid": "^8.3.2",
     "ws": "^7.4.6"
   },

commons/EventSubscriber.ts

@@ -0,0 +1,55 @@
+import IRegisteredEventListener, {
+  IEventSubscriber,
+} from '@secret-agent/interfaces/IRegisteredEventListener';
+
+export default class EventSubscriber implements IEventSubscriber {
+  private readonly registeredEventListeners: IRegisteredEventListener[] = [];
+
+  on<K extends string | symbol, Func extends (...args: any[]) => void>(
+    emitter: {
+      on(event: K, listener: Func, includeUnhandledEvents?: boolean);
+      off(event: K, listener: Func);
+    },
+    eventName: K,
+    handler: Func,
+    includeUnhandledEvents?: boolean,
+  ): IRegisteredEventListener {
+    emitter.on(eventName, handler, includeUnhandledEvents);
+    const registeredEvent: IRegisteredEventListener = { emitter, eventName, handler };
+    this.registeredEventListeners.push(registeredEvent);
+    return registeredEvent;
+  }
+
+  once<Event extends string | symbol, Func extends (...args: any[]) => void>(
+    emitter: {
+      once(event: Event, listener: Func, includeUnhandledEvents?: boolean);
+      off(event: Event, listener: Func);
+    },
+    eventName: Event,
+    handler: Func,
+    includeUnhandledEvents?: boolean,
+  ): IRegisteredEventListener {
+    emitter.once(eventName, handler, includeUnhandledEvents);
+    const registeredEvent: IRegisteredEventListener = { emitter, eventName, handler };
+    this.registeredEventListeners.push(registeredEvent);
+    return registeredEvent;
+  }
+
+  off(...listeners: IRegisteredEventListener[]): void {
+    for (const listener of listeners) {
+      listener.emitter.off(listener.eventName, listener.handler);
+    }
+    listeners.length = 0;
+  }
+
+  close(...keepMockEvents: (string | symbol)[]): void {
+    for (const listener of this.registeredEventListeners) {
+      if (keepMockEvents.includes(listener.eventName)) {
+        // add a mock event handler (like for capturing events)
+        (listener.emitter as any).on(listener.eventName, () => null);
+      }
+      listener.emitter.off(listener.eventName, listener.handler);
+    }
+    this.registeredEventListeners.length = 0;
+  }
+}

commons/Resolvable.ts

@@ -1,13 +1,12 @@
 import IResolvablePromise from '@secret-agent/interfaces/IResolvablePromise';
-import { bindFunctions } from './utils';
 import TimeoutError from './interfaces/TimeoutError';
 
 export default class Resolvable<T = any> implements IResolvablePromise<T>, PromiseLike<T> {
   public isResolved = false;
   public resolved: T;
   public promise: Promise<T>;
   public readonly timeout: NodeJS.Timeout;
-  public readonly stack: string;
+  public stack: string;
 
   private resolveFn: (value: T | PromiseLike<T>) => void;
   private rejectFn: (error?: Error) => void;
@@ -16,37 +15,42 @@ export default class Resolvable<T = any> implements IResolvablePromise<T>, Promi
     // get parent stack
     this.stack = new Error('').stack.slice(8);
 
+    this.promise = new Promise<T>((resolve, reject) => {
+      this.resolveFn = resolve;
+      this.rejectFn = reject;
+    });
+
     if (timeoutMillis !== undefined && timeoutMillis !== null) {
       this.timeout = setTimeout(
         this.rejectWithTimeout.bind(this, timeoutMessage),
         timeoutMillis,
       ).unref();
     }
-    this.promise = new Promise<T>((resolve, reject) => {
-      this.resolveFn = resolve;
-      this.rejectFn = reject;
-    });
-    bindFunctions(this);
+    this.resolve = this.resolve.bind(this);
+    this.reject = this.reject.bind(this);
   }
 
   public resolve(value: T | PromiseLike<T>): void {
     if (this.isResolved) return;
-    clearTimeout(this.timeout);
+
     this.resolveFn(value);
+
     Promise.resolve(value)
       // eslint-disable-next-line promise/always-return
       .then(x => {
         this.isResolved = true;
         this.resolved = x;
+        this.clean();
+        this.stack = null;
       })
       .catch(this.reject);
   }
 
   public reject(error: Error): void {
     if (this.isResolved) return;
     this.isResolved = true;
-    clearTimeout(this.timeout);
     this.rejectFn(error);
+    this.clean();
   }
 
   public toJSON(): object {
@@ -73,6 +77,12 @@ export default class Resolvable<T = any> implements IResolvablePromise<T>, Promi
     return this.promise.finally(onfinally);
   }
 
+  private clean(): void {
+    clearTimeout(this.timeout);
+    this.resolveFn = null;
+    this.rejectFn = null;
+  }
+
   private rejectWithTimeout(message: string): void {
     const error = new TimeoutError(message);
     error.stack = `TimeoutError: ${message}\n${this.stack}`;

commons/eventUtils.ts

@@ -5,58 +5,6 @@ import { IBoundLog } from '@secret-agent/interfaces/ILog';
 import { createPromise } from './utils';
 import IPendingWaitEvent, { CanceledPromiseError } from './interfaces/IPendingWaitEvent';
 
-export function addEventListener(
-  emitter: EventEmitter,
-  eventName: string | symbol,
-  handler: (...args: any[]) => void,
-): IRegisteredEventListener {
-  emitter.on(eventName, handler);
-  return { emitter, eventName, handler };
-}
-
-export function addEventListeners(
-  emitter: EventEmitter,
-  registrations: [string | symbol, (...args: any[]) => void, boolean?][],
-): IRegisteredEventListener[] {
-  return registrations.map(([eventName, handler]) => {
-    emitter.on(eventName, handler);
-    return { emitter, eventName, handler };
-  });
-}
-
-export function removeEventListeners(
-  listeners: Array<{
-    emitter: EventEmitter | ITypedEventEmitter<any>;
-    eventName: string | symbol;
-    handler: (...args: any[]) => void;
-  }>,
-): void {
-  for (const listener of listeners) {
-    listener.emitter.removeListener(listener.eventName, listener.handler);
-  }
-  listeners.length = 0;
-}
-
-export function addTypedEventListener<T, K extends keyof T & (string | symbol)>(
-  emitter: TypedEventEmitter<T>,
-  eventName: K,
-  handler: (this: TypedEventEmitter<T>, event?: T[K], initiator?: any) => any,
-  includeUnhandledEvents?: boolean,
-): IRegisteredEventListener {
-  emitter.on(eventName, handler, includeUnhandledEvents);
-  return { emitter, eventName, handler };
-}
-
-export function addTypedEventListeners<T, K extends keyof T & (string | symbol)>(
-  emitter: TypedEventEmitter<T>,
-  registrations: [K, (this: TypedEventEmitter<T>, event?: T[K]) => any, boolean?][],
-): IRegisteredEventListener[] {
-  return registrations.map(([eventName, handler, includeUnhandled]) => {
-    emitter.on(eventName, handler, includeUnhandled);
-    return { emitter, eventName, handler };
-  });
-}
-
 export class TypedEventEmitter<T> extends EventEmitter implements ITypedEventEmitter<T> {
   public storeEventsWithoutListeners = false;
 
@@ -110,18 +58,20 @@ export class TypedEventEmitter<T> extends EventEmitter implements ITypedEventEmi
       timeoutMillis,
     });
 
-    const listener = addTypedEventListener(this, eventType, (result: T[K]) => {
+    const callbackFn = (result: T[K]): void => {
       // give the listeners a second to register
       if (!listenerFn || listenerFn.call(this, result)) {
         this.logger?.stats(`waitOn.resolve:${eventType}`, {
           parentLogId: messageId,
         });
         promise.resolve(result);
       }
-    });
+    };
+
+    this.on(eventType, callbackFn);
 
     return promise.promise.finally(() => {
-      removeEventListeners([listener]);
+      this.off(eventType, callbackFn);
       const idx = this.pendingWaitEvents.findIndex(x => x.id === id);
       if (idx >= 0) this.pendingWaitEvents.splice(idx, 1);
     });

core/lib/AwaitedEventListener.ts

@@ -2,6 +2,7 @@ import { v1 as uuidv1 } from 'uuid';
 import { IJsPath } from 'awaited-dom/base/AwaitedPath';
 import ISessionMeta from '@secret-agent/interfaces/ISessionMeta';
 import { assert } from '@secret-agent/commons/utils';
+import EventSubscriber from '@secret-agent/commons/EventSubscriber';
 import IListenerObject from '../interfaces/IListenerObject';
 import Session from './Session';
 
@@ -14,9 +15,10 @@ export interface ITrigger {
 // TODO: can we merge TypedEventEmitter ideas and eventListenerIdsByType?. Revisit when we get to dom events
 export default class AwaitedEventListener {
   private readonly listenersById = new Map<string, IListenerObject>();
+  private eventSubscriber = new EventSubscriber();
 
-  constructor(readonly session: Session) {
-    session.on('closing', () => this.triggerListenersWithType('close'));
+  constructor(private session: Session) {
+    this.eventSubscriber.once(session, 'closing', () => this.triggerListenersWithType('close'));
   }
 
   public close() {
@@ -25,6 +27,8 @@ export default class AwaitedEventListener {
         this.remove(entry.id);
       }
     }
+    this.eventSubscriber.close();
+    this.session = null;
   }
 
   public listen(

core/lib/CommandRecorder.ts

@@ -12,8 +12,8 @@ export default class CommandRecorder {
   public readonly fnNames = new Set<string>();
   private logger: IBoundLog;
   constructor(
-    readonly owner: any,
-    readonly session: Session,
+    private owner: any,
+    private session: Session,
     readonly tabId: number,
     readonly frameId: number,
     fns: AsyncFunc[],
@@ -29,8 +29,14 @@ export default class CommandRecorder {
     });
   }
 
+  public clear(): void {
+    this.session = null;
+    this.owner = null;
+  }
+
   private async runCommandFn<T>(commandFn: AsyncFunc, ...args: any[]): Promise<T> {
-    if (!this.fnNames.has(commandFn.name)) throw new Error(`Unsupported function requested ${commandFn.name}`);
+    if (!this.fnNames.has(commandFn.name))
+      throw new Error(`Unsupported function requested ${commandFn.name}`);
 
     const { session } = this;
     const sessionState = session.sessionState;

core/lib/FrameEnvironment.ts

@@ -171,6 +171,7 @@ export default class FrameEnvironment {
       for (const path of this.cleanPaths) {
         Fs.promises.unlink(path).catch(() => null);
       }
+      this.commandRecorder.clear();
     } catch (error) {
       if (!error.message.includes('Target closed') && !(error instanceof CanceledPromiseError)) {
         this.logger.error('FrameEnvironment.ClosingError', { error, parentLogId });

core/lib/Session.ts

@@ -21,6 +21,7 @@ import IViewport from '@secret-agent/interfaces/IViewport';
 import IJsPathResult from '@secret-agent/interfaces/IJsPathResult';
 import ISessionCreateOptions from '@secret-agent/interfaces/ISessionCreateOptions';
 import IGeolocation from '@secret-agent/interfaces/IGeolocation';
+import EventSubscriber from '@secret-agent/commons/EventSubscriber';
 import SessionState from './SessionState';
 import AwaitedEventListener from './AwaitedEventListener';
 import GlobalPool from './GlobalPool';
@@ -80,6 +81,7 @@ export default class Session extends TypedEventEmitter<{
   private isolatedMitmProxy?: MitmProxy;
   private _isClosing = false;
   private detachedTabsById = new Map<number, Tab>();
+  private eventSubscriber = new EventSubscriber();
 
   private tabIdCounter = 0;
   private frameIdCounter = 0;
@@ -212,7 +214,7 @@ export default class Session extends TypedEventEmitter<{
       detachedState.detachedAtCommandId,
     );
     this.detachedTabsById.set(newTab.id, newTab);
-    newTab.on('close', () => {
+    newTab.once('close', () => {
       if (newTab.mainFrameEnvironment.jsPath.hasNewExecJsPathHistory) {
         this.sessionState.recordDetachedJsPathCalls(
           newTab.mainFrameEnvironment.jsPath.execHistory,
@@ -253,20 +255,22 @@ export default class Session extends TypedEventEmitter<{
 
   public async initialize(context: IPuppetContext) {
     this.browserContext = context;
-    context.on('devtools-message', this.onDevtoolsMessage.bind(this));
+
+    this.eventSubscriber.on(context, 'devtools-message', this.onDevtoolsMessage.bind(this));
+
     if (this.userProfile) {
       await UserProfile.install(this);
     }
 
     context.defaultPageInitializationFn = InjectedScripts.install;
 
     const requestSession = this.mitmRequestSession;
-    requestSession.on('request', this.onMitmRequest.bind(this));
-    requestSession.on('response', this.onMitmResponse.bind(this));
-    requestSession.on('http-error', this.onMitmError.bind(this));
-    requestSession.on('resource-state', this.onResourceStates.bind(this));
-    requestSession.on('socket-close', this.onSocketClose.bind(this));
-    requestSession.on('socket-connect', this.onSocketConnect.bind(this));
+    this.eventSubscriber.on(requestSession, 'request', this.onMitmRequest.bind(this));
+    this.eventSubscriber.on(requestSession, 'response', this.onMitmResponse.bind(this));
+    this.eventSubscriber.on(requestSession, 'http-error', this.onMitmError.bind(this));
+    this.eventSubscriber.on(requestSession, 'resource-state', this.onResourceStates.bind(this));
+    this.eventSubscriber.on(requestSession, 'socket-close', this.onSocketClose.bind(this));
+    this.eventSubscriber.on(requestSession, 'socket-connect', this.onSocketConnect.bind(this));
     await this.plugins.onHttpAgentInitialized(requestSession.requestAgent);
   }
 
@@ -332,6 +336,9 @@ export default class Session extends TypedEventEmitter<{
       sessionId: this.id,
       parentLogId: start,
     });
+    this.commandRecorder.clear();
+    this.eventSubscriber.close();
+
     this.emit('closed');
     // should go last so we can capture logs
     this.sessionState.close();

core/lib/SessionState.ts

@@ -56,7 +56,7 @@ export default class SessionState {
 
   private readonly logger: IBoundLog;
 
-  private readonly browserRequestIdToResources: {
+  private browserRequestIdToResources: {
     [browserRequestId: string]: { resourceId: number; url: string }[];
   } = {};
 
@@ -465,6 +465,10 @@ export default class SessionState {
     loggerSessionIdNames.delete(this.sessionId);
     this.db.flush();
     this.db.close();
+    this.resourcesById.clear();
+    this.browserRequestIdToResources = {};
+    this.websocketListeners = {};
+    this.websocketMessages.length = 0;
     SessionState.registry.delete(this.sessionId);
   }
 

core/lib/Tab.ts

@@ -252,6 +252,7 @@ export default class Tab extends TypedEventEmitter<ITabEventParams> {
         errors.push(error);
       }
     }
+    this.commandRecorder.clear();
     this.emit('close');
     this.logger.stats('Tab.Closed', { parentLogId, errors });
   }