Fix removing listeners in wrapped event emitters

The `.emitWithContext` function in the `ScopeManager`, which is
called from the `.wrapEmitter` function in the `Tracer`, modifies
a given `EventEmitter` so that the callbacks passed to it when
registering a listener are bound to the active span of the
asynchronous context in which the listener was registered. That is,
calling `scopeManager.active()` or `tracer.currentSpan()` from the
listener should return the span that was active at the point where
the listener was registered. This is done by wrapping the callback
in a function that closes over the current span, calling the
original callback from within `scopeManager.withContext` with the
closed over span, and registering this wrapper function as the
callback for the event emitter's listener.

Since the callback that is registered for an event emitter's
listener is a wrapper function over the callback, and not the
original callback, future attempts to remove the listener from the
event emitter silently fail, as the callback function that is
passed to `.removeListener` is not the callback function that was
originally registered.

This commit fixes this by keeping track of which wrapper functions
were created for which events and callbacks, and modifying the
given event emitter further, intercepting calls to methods that
remove listeners and replacing the given callback with its
previously created wrapper.

This commit also adds tests for the .emitWithContext function.

packages/nodejs/.changesets/fix-removing-event-listeners-from-wrapped-event-emitters.md

@@ -0,0 +1,8 @@
+---
+bump: "patch"
+type: "fix"
+---
+
+Fix removing event listeners from wrapped event emitters. When using
+`tracer.wrapEmitter` to wrap an event emitter, it was not possible to remove
+any listeners that were added after the event emitter was wrapped.

packages/nodejs/src/__tests__/scope.test.ts

@@ -1,6 +1,7 @@
 import { ScopeManager } from "../scope"
 import { RootSpan, ChildSpan } from "../span"
 import { Span } from "../interfaces/span"
+import { EventEmitter } from "events"
 
 describe("ScopeManager", () => {
   let scopeManager: ScopeManager
@@ -211,6 +212,83 @@ describe("ScopeManager", () => {
     })
   })
 
-  // TODO: Add tests
-  describe(".emitWithContext()", () => {})
+  describe(".emitWithContext()", () => {
+    let eventEmitter: EventEmitter
+    let listener: jest.Mock
+
+    beforeEach(() => {
+      eventEmitter = new EventEmitter()
+      listener = jest.fn(() => scopeManager.active())
+      scopeManager.emitWithContext(eventEmitter)
+    })
+
+    it("can add event listeners", () => {
+      eventEmitter.on("test", listener)
+
+      eventEmitter.emit("test")
+      expect(listener).toHaveBeenCalledTimes(1)
+    })
+
+    it("can remove event listeners", () => {
+      eventEmitter.on("test", listener)
+      eventEmitter.off("test", listener)
+
+      eventEmitter.emit("test")
+      expect(listener).not.toHaveBeenCalled()
+    })
+
+    it("binds event listeners to the active span of the scope that adds them", async () => {
+      const rootSpan = new RootSpan()
+      scopeManager.setRoot(rootSpan)
+
+      const listenerSpan = new ChildSpan(rootSpan)
+      await asyncTaskWithContext(scopeManager, listenerSpan, () => {
+        eventEmitter.on("test", listener)
+      })
+
+      eventEmitter.emit("test")
+      expect(listener).toHaveBeenCalledTimes(1)
+      expect(listener.mock.results[0].value).toBe(listenerSpan)
+    })
+
+    it("can bind the same event and callback to different active spans", async () => {
+      const firstSpan = new RootSpan()
+      await asyncTaskWithContext(scopeManager, firstSpan, () => {
+        eventEmitter.on("test", listener)
+      })
+
+      const secondSpan = new RootSpan()
+      await asyncTaskWithContext(scopeManager, secondSpan, () => {
+        eventEmitter.on("test", listener)
+      })
+
+      eventEmitter.emit("test")
+      expect(listener).toHaveBeenCalledTimes(2)
+      expect(listener.mock.results[0].value).toBe(firstSpan)
+      expect(listener.mock.results[1].value).toBe(secondSpan)
+    })
+
+    it("removes event listeners from before .emitWithContext was called", () => {
+      eventEmitter = new EventEmitter()
+      eventEmitter.on("test", listener)
+
+      scopeManager.emitWithContext(eventEmitter)
+      eventEmitter.on("test", listener)
+
+      eventEmitter.emit("test")
+      expect(listener).toHaveBeenCalledTimes(2)
+      listener.mockClear()
+
+      eventEmitter.off("test", listener)
+
+      eventEmitter.emit("test")
+      expect(listener).toHaveBeenCalledTimes(1)
+      listener.mockClear()
+
+      eventEmitter.off("test", listener)
+
+      eventEmitter.emit("test")
+      expect(listener).not.toHaveBeenCalled()
+    })
+  })
 })

packages/nodejs/src/scope.ts

@@ -15,14 +15,20 @@ import { EventEmitter } from "events"
 import shimmer from "shimmer"
 
 // A list of well-known EventEmitter methods that add event listeners.
-const EVENT_EMITTER_METHODS: Array<keyof EventEmitter> = [
+const EVENT_EMITTER_ADD_METHODS: Array<keyof EventEmitter> = [
   "addListener",
   "on",
   "once",
   "prependListener",
   "prependOnceListener"
 ]
 
+// A list of well-known EventEmitter methods that remove event listeners.
+const EVENT_EMITTER_REMOVE_METHODS: Array<keyof EventEmitter> = [
+  "off",
+  "removeListener"
+]
+
 const WRAPPED = Symbol("@appsignal/nodejs:WRAPPED")
 
 type ContextWrapped<T> = T & { [WRAPPED]?: boolean }
@@ -235,16 +241,75 @@ export class ScopeManager {
 
   public emitWithContext(ee: EventEmitter): void {
     // eslint-disable-next-line @typescript-eslint/no-this-alias
-    const that = this
+    const thisScopeManager = this
 
-    EVENT_EMITTER_METHODS.forEach(method => {
+    // For each event and callback that is registered with the event emitter,
+    // `boundCallbacks` keeps a reference to the newly-created context-bound
+    // callback that wraps around the original callback.
+    // A listener with the same event and callback could be added from
+    // different asynchronous contexts, meaning its context-bound wrapper
+    // should bind it to a different scope, so it keeps a list of
+    // context-bound callbacks for a single event and callback pair.
+    const boundCallbacks = new BoundCallbacks()
+
+    EVENT_EMITTER_ADD_METHODS.forEach(method => {
       if (ee[method]) {
         shimmer.wrap(ee, method, (oldFn: Func<any>) => {
           return function (this: unknown, event: string, cb: Func<void>) {
-            return oldFn.call(this, event, that.bindContext(cb))
+            const boundCallback = thisScopeManager.bindContext(cb)
+            boundCallbacks.push(event, cb, boundCallback)
+            return oldFn.call(this, event, boundCallback)
           }
         })
       }
     })
+
+    EVENT_EMITTER_REMOVE_METHODS.forEach(method => {
+      if (ee[method]) {
+        shimmer.wrap(ee, method, (oldFn: Func<any>) => {
+          return function (this: unknown, event: string, cb: Func<void>) {
+            // If there is no bound callback for this event and callback, it
+            // might be a listener that was added before the event emitter was
+            // wrapped, so we should attempt to remove the given callback.
+            const maybeBoundCallback = boundCallbacks.pop(event, cb) ?? cb
+            return oldFn.call(this, event, maybeBoundCallback)
+          }
+        })
+      }
+    })
+  }
+}
+
+class BoundCallbacks {
+  #map: Map<string, WeakMap<Func<void>, Func<void>[]>>
+
+  constructor() {
+    this.#map = new Map()
+  }
+
+  push(event: string, cb: Func<void>, boundCallback: Func<void>): void {
+    let eventMap = this.#map.get(event)
+    if (!eventMap) {
+      eventMap = new WeakMap()
+      this.#map.set(event, eventMap)
+    }
+
+    let callbacks = eventMap.get(cb)
+    if (!callbacks) {
+      callbacks = []
+      eventMap.set(cb, callbacks)
+    }
+
+    callbacks.push(boundCallback)
+  }
+
+  pop(event: string, cb: Func<void>): Func<void> | undefined {
+    const eventMap = this.#map.get(event)
+    if (!eventMap) return
+
+    const callbacks = eventMap.get(cb)
+    if (!callbacks) return
+
+    return callbacks.pop()
   }
 }