From 5ca65e06c3defd2fb8c96422dca54716f7962174 Mon Sep 17 00:00:00 2001
From: jrandolf <101637635+jrandolf@users.noreply.github.com>
Date: Tue, 23 Jan 2024 16:08:20 +0100
Subject: [PATCH] chore: implement higher order event emitters (#11723)

---
 .../src/common/EventEmitter.test.ts           |  22 ++++
 .../puppeteer-core/src/common/EventEmitter.ts | 108 +++++++++++-------
 2 files changed, 91 insertions(+), 39 deletions(-)

diff --git a/packages/puppeteer-core/src/common/EventEmitter.test.ts b/packages/puppeteer-core/src/common/EventEmitter.test.ts
index c6b58bd4da7..cf05ef67001 100644
--- a/packages/puppeteer-core/src/common/EventEmitter.test.ts
+++ b/packages/puppeteer-core/src/common/EventEmitter.test.ts
@@ -160,4 +160,26 @@ describe('EventEmitter', () => {
       expect(emitter.emit('bar', undefined)).toBe(false);
     });
   });
+
+  describe('dispose', () => {
+    it('should dispose higher order emitters properly', () => {
+      let values = '';
+      emitter.on('foo', () => {
+        values += '1';
+      });
+      const higherOrderEmitter = new EventEmitter(emitter);
+
+      higherOrderEmitter.on('foo', () => {
+        values += '2';
+      });
+      higherOrderEmitter.emit('foo', undefined);
+
+      expect(values).toMatch('12');
+
+      higherOrderEmitter.off('foo');
+      higherOrderEmitter.emit('foo', undefined);
+
+      expect(values).toMatch('121');
+    });
+  });
 });
diff --git a/packages/puppeteer-core/src/common/EventEmitter.ts b/packages/puppeteer-core/src/common/EventEmitter.ts
index 9a39cc0a848..4a8bcb801f7 100644
--- a/packages/puppeteer-core/src/common/EventEmitter.ts
+++ b/packages/puppeteer-core/src/common/EventEmitter.ts
@@ -4,10 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import mitt, {
-  type Emitter,
-  type EventHandlerMap,
-} from '../../third_party/mitt/mitt.js';
+import mitt, {type Emitter} from '../../third_party/mitt/mitt.js';
 import {disposeSymbol} from '../util/disposable.js';
 
 /**
@@ -74,14 +71,20 @@ export type EventsWithWildcard<Events extends Record<EventType, unknown>> =
 export class EventEmitter<Events extends Record<EventType, unknown>>
   implements CommonEventEmitter<EventsWithWildcard<Events>>
 {
-  #emitter: Emitter<Events & {'*': Events[keyof Events]}>;
-  #handlers: EventHandlerMap<Events & {'*': Events[keyof Events]}> = new Map();
+  #emitter: Emitter<EventsWithWildcard<Events>> | EventEmitter<Events>;
+  #handlers = new Map<keyof Events | '*', Array<Handler<any>>>();
 
   /**
+   * If you pass an emitter, the returned emitter will wrap the passed emitter.
+   *
    * @internal
    */
-  constructor() {
-    this.#emitter = mitt(this.#handlers);
+  constructor(
+    emitter: Emitter<EventsWithWildcard<Events>> | EventEmitter<Events> = mitt(
+      new Map()
+    )
+  ) {
+    this.#emitter = emitter;
   }
 
   /**
@@ -94,6 +97,13 @@ export class EventEmitter<Events extends Record<EventType, unknown>>
     type: Key,
     handler: Handler<EventsWithWildcard<Events>[Key]>
   ): this {
+    const handlers = this.#handlers.get(type);
+    if (handlers === undefined) {
+      this.#handlers.set(type, [handler]);
+    } else {
+      handlers.push(handler);
+    }
+
     this.#emitter.on(type, handler);
     return this;
   }
@@ -108,33 +118,18 @@ export class EventEmitter<Events extends Record<EventType, unknown>>
     type: Key,
     handler?: Handler<EventsWithWildcard<Events>[Key]>
   ): this {
-    this.#emitter.off(type, handler);
-    return this;
-  }
-
-  /**
-   * Remove an event listener.
-   *
-   * @deprecated please use {@link EventEmitter.off} instead.
-   */
-  removeListener<Key extends keyof EventsWithWildcard<Events>>(
-    type: Key,
-    handler: Handler<EventsWithWildcard<Events>[Key]>
-  ): this {
-    this.off(type, handler);
-    return this;
-  }
-
-  /**
-   * Add an event listener.
-   *
-   * @deprecated please use {@link EventEmitter.on} instead.
-   */
-  addListener<Key extends keyof EventsWithWildcard<Events>>(
-    type: Key,
-    handler: Handler<EventsWithWildcard<Events>[Key]>
-  ): this {
-    this.on(type, handler);
+    const handlers = this.#handlers.get(type) ?? [];
+    if (handler === undefined) {
+      for (const handler of handlers) {
+        this.#emitter.off(type, handler);
+      }
+      this.#handlers.delete(type);
+      return this;
+    }
+    const index = handlers.lastIndexOf(handler);
+    if (index > -1) {
+      this.#emitter.off(type, ...handlers.splice(index, 1));
+    }
     return this;
   }
 
@@ -153,6 +148,30 @@ export class EventEmitter<Events extends Record<EventType, unknown>>
     return this.listenerCount(type) > 0;
   }
 
+  /**
+   * Remove an event listener.
+   *
+   * @deprecated please use {@link EventEmitter.off} instead.
+   */
+  removeListener<Key extends keyof EventsWithWildcard<Events>>(
+    type: Key,
+    handler: Handler<EventsWithWildcard<Events>[Key]>
+  ): this {
+    return this.off(type, handler);
+  }
+
+  /**
+   * Add an event listener.
+   *
+   * @deprecated please use {@link EventEmitter.on} instead.
+   */
+  addListener<Key extends keyof EventsWithWildcard<Events>>(
+    type: Key,
+    handler: Handler<EventsWithWildcard<Events>[Key]>
+  ): this {
+    return this.on(type, handler);
+  }
+
   /**
    * Like `on` but the listener will only be fired once and then it will be removed.
    * @param type - the event you'd like to listen to
@@ -189,13 +208,24 @@ export class EventEmitter<Events extends Record<EventType, unknown>>
    * @returns `this` to enable you to chain method calls.
    */
   removeAllListeners(type?: keyof EventsWithWildcard<Events>): this {
-    if (type === undefined || type === '*') {
-      this.#handlers.clear();
-    } else {
-      this.#handlers.delete(type);
+    if (type !== undefined) {
+      return this.off(type);
     }
+    this[disposeSymbol]();
     return this;
   }
+
+  /**
+   * @internal
+   */
+  [disposeSymbol](): void {
+    for (const [type, handlers] of this.#handlers) {
+      for (const handler of handlers) {
+        this.#emitter.off(type, handler);
+      }
+    }
+    this.#handlers.clear();
+  }
 }
 
 /**