Merge dd9e1bf into 4a3fd1f

open-telemetry · Apr 9, 2021 · 15f5dc1 · 15f5dc1
2 parents 4a3fd1f + dd9e1bf
commit 15f5dc1
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 deletions.
diff --git a/packages/opentelemetry-context-async-hooks/src/AbstractAsyncHooksContextManager.ts b/packages/opentelemetry-context-async-hooks/src/AbstractAsyncHooksContextManager.ts
@@ -146,10 +146,14 @@ export abstract class AbstractAsyncHooksContextManager
     const contextManager = this;
     return function (this: never, event: string) {
       const map = contextManager._getPatchMap(ee);
-      if (map?.[event] !== undefined) {
-        delete map[event];
+      if (map !== undefined) {
+        if (arguments.length === 0) {
+          contextManager._createPatchMap(ee);
+        } else if (map[event] !== undefined) {
+          delete map[event];
+        }
       }
-      return original.call(this, event);
+      return original.apply(this, arguments);
     };
   }
 
@@ -184,7 +188,7 @@ export abstract class AbstractAsyncHooksContextManager
   }
 
   private _createPatchMap(ee: EventEmitter): PatchMap {
-    const map = {};
+    const map = Object.create(null);
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     (ee as any)[this._kOtListeners] = map;
     return map;

diff --git a/packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts b/packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts
@@ -415,6 +415,24 @@ for (const contextManagerClass of [
         patchedEE.emit('test');
       });
 
+      it('should return current context and removeAllListeners (when enabled)', done => {
+        const ee = new EventEmitter();
+        const context = ROOT_CONTEXT.setValue(key1, 1);
+        const patchedEE = contextManager.bind(ee, context);
+        const handler = () => {
+          assert.deepStrictEqual(contextManager.active(), context);
+          patchedEE.removeAllListeners();
+          assert.strictEqual(patchedEE.listeners('test').length, 0);
+          assert.strictEqual(patchedEE.listeners('test1').length, 0);
+          return done();
+        };
+        patchedEE.on('test', handler);
+        patchedEE.on('test1', handler);
+        assert.strictEqual(patchedEE.listeners('test').length, 1);
+        assert.strictEqual(patchedEE.listeners('test1').length, 1);
+        patchedEE.emit('test');
+      });
+
       /**
        * Even if asynchooks is disabled, the context propagation will
        * still works but it might be lost after any async op.