diff --git a/src/util.js b/src/util.js index 89a3845..a268e69 100644 --- a/src/util.js +++ b/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) + }) } } diff --git a/tests/util.test.js b/tests/util.test.js index 2d8de4c..233bbfb 100644 --- a/tests/util.test.js +++ b/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) + }) })