diff --git a/packages/opentelemetry-scope-async-hooks/src/AsyncHooksScopeManager.ts b/packages/opentelemetry-scope-async-hooks/src/AsyncHooksScopeManager.ts index e6144811c3e..dc3ed0335e2 100644 --- a/packages/opentelemetry-scope-async-hooks/src/AsyncHooksScopeManager.ts +++ b/packages/opentelemetry-scope-async-hooks/src/AsyncHooksScopeManager.ts @@ -80,6 +80,8 @@ export class AsyncHooksScopeManager implements ScopeManager { return this._bindEventEmitter(target, scope); } else if (typeof target === 'function') { return this._bindFunction(target, scope); + } else if (target instanceof Promise) { + return this._bindPromise(target, scope); } return target; } @@ -145,6 +147,28 @@ export class AsyncHooksScopeManager implements ScopeManager { return target; } + private _bindPromise>( + target: T, + scope?: unknown + ): T { + const scopeManager = this; + const then = target.then; + target.then = function(this: {}) { + const args = new Array(arguments.length); + for (let i = 0; i < args.length; i++) { + args[i] = scopeManager.bind(arguments[i], scope); + } + /** + * Forced to remove compilation error here, the type of Promise + * is impossible(?) to reproduce when modifying it. + */ + // tslint:disable-next-line:ban-ts-ignore + // @ts-ignore + return then.apply(this, args); + }; + return target; + } + /** * Patch methods that remove a given listener so that we match the "patched" * version of that listener (the one that propagate context). diff --git a/packages/opentelemetry-scope-async-hooks/test/asynchooks/AsyncHooksScopeManager.test.ts b/packages/opentelemetry-scope-async-hooks/test/asynchooks/AsyncHooksScopeManager.test.ts index 4ade6914b51..e3fdac8c72a 100644 --- a/packages/opentelemetry-scope-async-hooks/test/asynchooks/AsyncHooksScopeManager.test.ts +++ b/packages/opentelemetry-scope-async-hooks/test/asynchooks/AsyncHooksScopeManager.test.ts @@ -231,4 +231,60 @@ describe('AsyncHooksScopeManager', () => { patchedEe.emit('test'); }); }); + + describe('.bind(promise)', () => { + it('should return the same target (when enabled)', () => { + const promise = Promise.resolve(); + assert.deepStrictEqual(scopeManager.bind(promise), promise); + }); + + it('should return the scope (when enabled)', done => { + const scope = { a: 3 }; + const promise = Promise.resolve(); + scopeManager.bind(promise, scope); + promise.then(() => { + assert.deepStrictEqual(scopeManager.active(), scope); + return done(); + }); + }); + + it('should return the scope (when disabled)', done => { + scopeManager.disable(); + const scope = { a: 3 }; + const promise = Promise.resolve(); + scopeManager.bind(promise, scope); + promise.then(() => { + assert.deepStrictEqual(scopeManager.active(), scope); + scopeManager.enable(); + return done(); + }); + }); + + it('should not return the scope (when disabled + async op)', done => { + scopeManager.disable(); + const scope = { a: 3 }; + const promise = Promise.resolve(); + scopeManager.bind(promise, scope); + promise.then(() => { + setImmediate(() => { + assert.deepStrictEqual(scopeManager.active(), null); + scopeManager.enable(); + return done(); + }); + }); + }); + + it('should return the scope (when enabled + async op)', done => { + const scope = { a: 3 }; + const promise = Promise.resolve(); + scopeManager.bind(promise, scope); + promise.then(() => { + setImmediate(() => { + assert.deepStrictEqual(scopeManager.active(), scope); + scopeManager.enable(); + return done(); + }); + }); + }); + }); });