fix: event-handler leakage (#788)

Fixes an issue where event handlers could be leaked if the same handler
reference was added multiple times.

---------

Signed-off-by: Todd Baert <[email protected]>

Author: toddbaert

packages/shared/src/events/generic-event-emitter.ts

@@ -1,6 +1,6 @@
 import { Logger, ManageLogger, SafeLogger } from '../logger';
 import { EventContext, EventDetails, EventHandler } from './eventing';
-import { AnyProviderEvent } from './events';
+import { AllProviderEvents, AnyProviderEvent } from './events';
 
 /**
  * The GenericEventEmitter should only be used within the SDK. It supports additional properties that can be included
@@ -11,8 +11,13 @@ export abstract class GenericEventEmitter<E extends AnyProviderEvent, Additional
 {
   protected abstract readonly eventEmitter: PlatformEventEmitter;
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  private readonly _handlers = new WeakMap<EventHandler, EventHandler>();
+  private readonly _handlers: { [key in AnyProviderEvent]: WeakMap<EventHandler, EventHandler[]>} = {
+    [AllProviderEvents.ConfigurationChanged]: new WeakMap<EventHandler, EventHandler[]>(),
+    [AllProviderEvents.ContextChanged]: new WeakMap<EventHandler, EventHandler[]>(),
+    [AllProviderEvents.Ready]: new WeakMap<EventHandler, EventHandler[]>(),
+    [AllProviderEvents.Error]: new WeakMap<EventHandler, EventHandler[]>(),
+    [AllProviderEvents.Stale]: new WeakMap<EventHandler, EventHandler[]>(),
+  };
   private _eventLogger?: Logger;
 
   constructor(private readonly globalLogger?: () => Logger) {}
@@ -29,19 +34,25 @@ export abstract class GenericEventEmitter<E extends AnyProviderEvent, Additional
       await handler(details);    
     };
     // The async handler has to be written to the map, because we need to get the wrapper function when deleting a listener
-    this._handlers.set(handler, asyncHandler);
+    const existingAsyncHandlers = this._handlers[eventType].get(handler);
+
+    // we allow duplicate event handlers, similar to node,
+    // see: https://nodejs.org/api/events.html#emitteroneventname-listener
+    // and https://nodejs.org/api/events.html#emitterremovelistenereventname-listener
+    this._handlers[eventType].set(handler, [...(existingAsyncHandlers || []), asyncHandler]);
     this.eventEmitter.on(eventType, asyncHandler);
   }
 
   removeHandler(eventType: AnyProviderEvent, handler: EventHandler): void {
     // Get the wrapper function for this handler, to delete it from the event emitter
-    const asyncHandler = this._handlers.get(handler) as EventHandler | undefined;
+    const existingAsyncHandlers = this._handlers[eventType].get(handler);
 
-    if (!asyncHandler) {
-      return;
+    if (existingAsyncHandlers) {
+      const removedAsyncHandler = existingAsyncHandlers.pop();
+      if (removedAsyncHandler) {
+        this.eventEmitter.removeListener(eventType, removedAsyncHandler);
+      }
     }
-
-    this.eventEmitter.removeListener(eventType, asyncHandler);
   }
 
   removeAllHandlers(eventType?: AnyProviderEvent): void {

packages/shared/test/events.spec.ts

@@ -13,19 +13,27 @@ class TestEventEmitter extends GenericEventEmitter<AnyProviderEvent> {
   }
 }
 
+// a little function to make sure we're at least waiting for the event loop 
+// to clear before we start making assertions
+const wait = (millis = 0) => {
+  return new Promise(resolve => {setTimeout(resolve, millis);});
+};
+
 describe('GenericEventEmitter', () => {
   describe('addHandler should', function () {
-    it('attach handler for event type', function () {
+    it('attach handler for event type', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
       emitter.addHandler(AllProviderEvents.Ready, handler1);
       emitter.emit(AllProviderEvents.Ready);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
     });
 
-    it('attach several handlers for event type', function () {
+    it('attach several handlers for event type', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
@@ -38,6 +46,8 @@ describe('GenericEventEmitter', () => {
 
       emitter.emit(AllProviderEvents.Ready);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
       expect(handler2).toHaveBeenCalledTimes(1);
       expect(handler3).not.toHaveBeenCalled();
@@ -62,28 +72,32 @@ describe('GenericEventEmitter', () => {
       emitter.emit(AllProviderEvents.Ready);
     });
 
-    it('trigger handler for event type', function () {
+    it('trigger handler for event type', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
       emitter.addHandler(AllProviderEvents.Ready, handler1);
       emitter.emit(AllProviderEvents.Ready);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
     });
 
-    it('trigger handler for event type with event data', function () {
+    it('trigger handler for event type with event data', async function () {
       const event: ReadyEvent = { message: 'message' };
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
       emitter.addHandler(AllProviderEvents.Ready, handler1);
       emitter.emit(AllProviderEvents.Ready, event);
 
+      await wait();
+
       expect(handler1).toHaveBeenNthCalledWith(1, event);
     });
 
-    it('trigger several handlers for event type', function () {
+    it('trigger several handlers for event type', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
@@ -96,14 +110,16 @@ describe('GenericEventEmitter', () => {
 
       emitter.emit(AllProviderEvents.Ready);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
       expect(handler2).toHaveBeenCalledTimes(1);
       expect(handler3).not.toHaveBeenCalled();
     });
   });
 
   describe('removeHandler should', () => {
-    it('remove single handler', function () {
+    it('remove single handler', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
@@ -113,12 +129,14 @@ describe('GenericEventEmitter', () => {
       emitter.removeHandler(AllProviderEvents.Ready, handler1);
       emitter.emit(AllProviderEvents.Ready);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
     });
   });
 
   describe('removeAllHandlers should', () => {
-    it('remove all handlers for event type', function () {
+    it('remove all handlers for event type', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
@@ -130,11 +148,62 @@ describe('GenericEventEmitter', () => {
       emitter.emit(AllProviderEvents.Ready);
       emitter.emit(AllProviderEvents.Error);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(0);
       expect(handler2).toHaveBeenCalledTimes(1);
     });
 
-    it('remove all handlers only for event type', function () {
+    it('remove same handler when assigned to multiple events', async function () {
+      const emitter = new TestEventEmitter();
+
+      const handler = jest.fn();
+      emitter.addHandler(AllProviderEvents.Stale, handler);
+      emitter.addHandler(AllProviderEvents.ContextChanged, handler);
+
+      emitter.removeHandler(AllProviderEvents.Stale, handler);
+      emitter.removeHandler(AllProviderEvents.ContextChanged, handler);
+
+      emitter.emit(AllProviderEvents.Stale);
+      emitter.emit(AllProviderEvents.ContextChanged);
+
+      await wait();
+
+      expect(handler).toHaveBeenCalledTimes(0);
+    });
+
+    it('allow addition/removal of duplicate handlers', async function () {
+      const emitter = new TestEventEmitter();
+
+      const handler = jest.fn();
+      emitter.addHandler(AllProviderEvents.Stale, handler);
+      emitter.addHandler(AllProviderEvents.Stale, handler);
+
+      emitter.removeHandler(AllProviderEvents.Stale, handler);
+      emitter.removeHandler(AllProviderEvents.Stale, handler);
+
+      emitter.emit(AllProviderEvents.Stale);
+
+      await wait();
+
+      expect(handler).toHaveBeenCalledTimes(0);
+    });
+
+    it('allow duplicate event handlers and call them', async function () {
+      const emitter = new TestEventEmitter();
+
+      const handler = jest.fn();
+      emitter.addHandler(AllProviderEvents.Stale, handler);
+      emitter.addHandler(AllProviderEvents.Stale, handler);
+
+      emitter.emit(AllProviderEvents.Stale);
+
+      await wait();
+
+      expect(handler).toHaveBeenCalledTimes(2);
+    });
+
+    it('remove all handlers only for event type', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
@@ -146,11 +215,13 @@ describe('GenericEventEmitter', () => {
       emitter.removeAllHandlers(AllProviderEvents.Ready);
       emitter.emit(AllProviderEvents.Ready);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
       expect(handler2).toHaveBeenCalledTimes(0);
     });
 
-    it('remove all handlers if no event type is given', function () {
+    it('remove all handlers if no event type is given', async function () {
       const emitter = new TestEventEmitter();
 
       const handler1 = jest.fn();
@@ -164,6 +235,8 @@ describe('GenericEventEmitter', () => {
       emitter.emit(AllProviderEvents.Ready);
       emitter.emit(AllProviderEvents.Error);
 
+      await wait();
+
       expect(handler1).toHaveBeenCalledTimes(1);
       expect(handler2).toHaveBeenCalledTimes(1);
     });