Prevent double wrapping event emitter listeners (open-telemetry#3133)

CHANGELOG.md

@@ -13,6 +13,9 @@ All notable changes to this project will be documented in this file.
 
 ### :bug: (Bug Fix)
 
+* fix(context-async-hooks): Ensure listeners added using `once` can be removed using `removeListener`
+  [#3133](https://github.com/open-telemetry/opentelemetry-js/pull/3133)
+
 ### :books: (Refine Doc)
 
 ### :house: (Internal)

packages/opentelemetry-context-async-hooks/src/AbstractAsyncHooksContextManager.ts

@@ -177,6 +177,17 @@ implements ContextManager {
   ) {
     const contextManager = this;
     return function (this: never, event: string, listener: Func<void>) {
+      /**
+       * This check is required to prevent double-wrapping the listener.
+       * The implementation for ee.once wraps the listener and calls ee.on.
+       * Without this check, we would wrap that wrapped listener.
+       * This causes an issue because ee.removeListener depends on the onceWrapper
+       * to properly remove the listener. If we wrap their wrapper, we break
+       * that detection.
+       */
+      if (contextManager._wrapped) {
+        return original.call(this, event, listener);
+      }
       let map = contextManager._getPatchMap(ee);
       if (map === undefined) {
         map = contextManager._createPatchMap(ee);
@@ -189,7 +200,16 @@ implements ContextManager {
       const patchedListener = contextManager.bind(context, listener);
       // store a weak reference of the user listener to ours
       listeners.set(listener, patchedListener);
-      return original.call(this, event, patchedListener);
+
+      /**
+       * See comment at the start of this function for the explanation of this property.
+       */
+      contextManager._wrapped = true;
+      try {
+        return original.call(this, event, patchedListener);
+      } finally {
+        contextManager._wrapped = false;
+      }
     };
   }
 
@@ -204,4 +224,5 @@ implements ContextManager {
   }
 
   private readonly _kOtListeners = Symbol('OtListeners');
+  private _wrapped = false;
 }

packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts

@@ -415,6 +415,28 @@ for (const contextManagerClass of [
         patchedEE.emit('test');
       });
 
+      it('should remove event handler enabled by .once using removeListener (when enabled)', () => {
+        const ee = new EventEmitter();
+        const context = ROOT_CONTEXT.setValue(key1, 1);
+        const patchedEE = contextManager.bind(context, ee);
+        function handler() {}
+        patchedEE.once('test', handler);
+        assert.strictEqual(patchedEE.listeners('test').length, 1);
+        patchedEE.removeListener('test', handler);
+        assert.strictEqual(patchedEE.listeners('test').length, 0);
+      });
+
+      it('should remove event handler enabled by .once using off (when enabled)', () => {
+        const ee = new EventEmitter();
+        const context = ROOT_CONTEXT.setValue(key1, 1);
+        const patchedEE = contextManager.bind(context, ee);
+        const handler = () => { };
+        patchedEE.once('test', handler);
+        assert.strictEqual(patchedEE.listeners('test').length, 1);
+        patchedEE.off('test', handler);
+        assert.strictEqual(patchedEE.listeners('test').length, 0);
+      });
+
       it('should return current context and removeAllListeners (when enabled)', done => {
         const ee = new EventEmitter();
         const context = ROOT_CONTEXT.setValue(key1, 1);