Fix EE wrapper to patch listeners instead of emit (#39)

src/util.js

@@ -1,14 +1,49 @@
 'use strict'
 
+const isWrappedSymbol = Symbol('cls-rtracer-is-wrapped')
+const wrappedSymbol = Symbol('cls-rtracer-wrapped-function')
+
+function wrapEmitterMethod (emitter, method, wrapper) {
+  if (emitter[method][isWrappedSymbol]) {
+    return
+  }
+
+  const original = emitter[method]
+  const wrapped = wrapper(original, method)
+  wrapped[isWrappedSymbol] = true
+  emitter[method] = wrapped
+
+  return wrapped
+}
+
+const addMethods = [
+  'on',
+  'addListener',
+  'prependListener'
+]
+
+const removeMethods = [
+  'off',
+  'removeListener'
+]
+
 /**
- * Monkey patches `.emit()` method of the given emitter, so
- * that all event listeners are run in scope of the provided
- * async resource.
+ * Wraps EventEmitter listener registration methods of the
+ * given emitter, so that all listeners are run in scope of
+ * the provided async resource.
  */
-const wrapEmitter = (emitter, asyncResource) => {
-  const original = emitter.emit
-  emitter.emit = function (type, ...args) {
-    return asyncResource.runInAsyncScope(original, emitter, type, ...args)
+function wrapEmitter (emitter, asyncResource) {
+  for (const method of addMethods) {
+    wrapEmitterMethod(emitter, method, (original) => function (name, handler) {
+      handler[wrappedSymbol] = asyncResource.runInAsyncScope.bind(asyncResource, handler, emitter)
+      return original.call(this, name, handler[wrappedSymbol])
+    })
+  }
+
+  for (const method of removeMethods) {
+    wrapEmitterMethod(emitter, method, (original) => function (name, handler) {
+      return original.call(this, name, handler[wrappedSymbol] || handler)
+    })
   }
 }
 

tests/util.test.js

@@ -6,11 +6,10 @@ const { AsyncResource, executionAsyncId } = require('async_hooks')
 const { wrapEmitter } = require('../src/util')
 
 describe('wrapEmitter', () => {
-  test('binds event listenters with async resource', (done) => {
+  test('binds event listeners with async resource', (done) => {
     const emitter = new EventEmitter()
     const asyncResource = new AsyncResource('foobar')
     wrapEmitter(emitter, asyncResource)
-
     setTimeout(() => {
       emitter.emit('foo')
     }, 0)
@@ -24,4 +23,52 @@ describe('wrapEmitter', () => {
       }
     })
   })
+
+  test('does not bind previously registered event listeners', (done) => {
+    const emitter = new EventEmitter()
+    const asyncResource = new AsyncResource('foobar')
+
+    emitter.on('foo', () => {
+      try {
+        expect(executionAsyncId()).not.toEqual(asyncResource.asyncId())
+        done()
+      } catch (error) {
+        done(error)
+      }
+    })
+
+    wrapEmitter(emitter, asyncResource)
+    setTimeout(() => {
+      emitter.emit('foo')
+    }, 0)
+  })
+
+  test('does not prevent event listeners from being removed', (done) => {
+    const emitter = new EventEmitter()
+    const asyncResource = new AsyncResource('foobar')
+    wrapEmitter(emitter, asyncResource)
+
+    const listener = () => {
+      done(new Error('Boom'))
+    }
+    emitter.on('foo', listener)
+    emitter.off('foo', listener)
+
+    emitter.emit('foo')
+    done()
+  })
+
+  test('wraps only once on multiple invocations', () => {
+    const emitter = new EventEmitter()
+    const asyncResource = new AsyncResource('foobar')
+
+    const unwrappedMethod = emitter.addListener
+    wrapEmitter(emitter, asyncResource)
+    const wrappedMethod1 = emitter.addListener
+    wrapEmitter(emitter, asyncResource)
+    const wrappedMethod2 = emitter.addListener
+
+    expect(unwrappedMethod).not.toEqual(wrappedMethod1)
+    expect(wrappedMethod1).toEqual(wrappedMethod2)
+  })
 })