From 4f9e8f2d09ddc20fde651380ca1344a8dd6e5085 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20W=C3=BCrbach?= Date: Wed, 17 Jun 2020 22:10:49 +0200 Subject: [PATCH] feat: improve ContextManager performance using AsyncLocalStorage Use AsyncLocalStorage in ContentManager to improve performance --- .../src/AsyncHooksContextManager.ts | 164 +---- .../src/AsyncLocalStorageContextManager.ts | 48 ++ .../src/BaseContextManager.ts | 188 ++++++ .../src/index.ts | 1 + .../test/AsyncHooksContextManager.test.ts | 559 +++++++++--------- 5 files changed, 534 insertions(+), 426 deletions(-) create mode 100644 packages/opentelemetry-context-async-hooks/src/AsyncLocalStorageContextManager.ts create mode 100644 packages/opentelemetry-context-async-hooks/src/BaseContextManager.ts diff --git a/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts b/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts index 0d22c1fd5e..b3eede09b9 100644 --- a/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts +++ b/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts @@ -14,35 +14,17 @@ * limitations under the License. */ -import { ContextManager, Context } from '@opentelemetry/context-base'; +import { Context } from '@opentelemetry/context-base'; import * as asyncHooks from 'async_hooks'; -import { EventEmitter } from 'events'; +import { BaseContextManager } from './BaseContextManager'; -type Func = (...args: unknown[]) => T; - -type PatchedEventEmitter = { - /** - * Store a map for each event of all original listener and their "patched" - * version so when the listener is removed by the user, we remove the - * corresponding "patched" function. - */ - __ot_listeners?: { [name: string]: WeakMap, Func> }; -} & EventEmitter; - -const ADD_LISTENER_METHODS = [ - 'addListener' as 'addListener', - 'on' as 'on', - 'once' as 'once', - 'prependListener' as 'prependListener', - 'prependOnceListener' as 'prependOnceListener', -]; - -export class AsyncHooksContextManager implements ContextManager { +export class AsyncHooksContextManager extends BaseContextManager { private _asyncHook: asyncHooks.AsyncHook; private _contexts: Map = new Map(); private _stack: Array = []; constructor() { + super(); this._asyncHook = asyncHooks.createHook({ init: this._init.bind(this), before: this._before.bind(this), @@ -68,19 +50,6 @@ export class AsyncHooksContextManager implements ContextManager { } } - bind(target: T, context?: Context): T { - // if no specific context to propagate is given, we use the current one - if (context === undefined) { - context = this.active(); - } - if (target instanceof EventEmitter) { - return this._bindEventEmitter(target, context); - } else if (typeof target === 'function') { - return this._bindFunction(target, context); - } - return target; - } - enable(): this { this._asyncHook.enable(); return this; @@ -93,131 +62,6 @@ export class AsyncHooksContextManager implements ContextManager { return this; } - private _bindFunction(target: T, context: Context): T { - const manager = this; - const contextWrapper = function (this: {}, ...args: unknown[]) { - return manager.with(context, () => target.apply(this, args)); - }; - Object.defineProperty(contextWrapper, 'length', { - enumerable: false, - configurable: true, - writable: false, - value: target.length, - }); - /** - * It isn't possible to tell Typescript that contextWrapper is the same as T - * so we forced to cast as any here. - */ - // eslint-disable-next-line @typescript-eslint/no-explicit-any - return contextWrapper as any; - } - - /** - * By default, EventEmitter call their callback with their context, which we do - * not want, instead we will bind a specific context to all callbacks that - * go through it. - * @param target EventEmitter a instance of EventEmitter to patch - * @param context the context we want to bind - */ - private _bindEventEmitter( - target: T, - context: Context - ): T { - const ee = (target as unknown) as PatchedEventEmitter; - if (ee.__ot_listeners !== undefined) return target; - ee.__ot_listeners = {}; - - // patch methods that add a listener to propagate context - ADD_LISTENER_METHODS.forEach(methodName => { - if (ee[methodName] === undefined) return; - ee[methodName] = this._patchAddListener(ee, ee[methodName], context); - }); - // patch methods that remove a listener - if (typeof ee.removeListener === 'function') { - ee.removeListener = this._patchRemoveListener(ee, ee.removeListener); - } - if (typeof ee.off === 'function') { - ee.off = this._patchRemoveListener(ee, ee.off); - } - // patch method that remove all listeners - if (typeof ee.removeAllListeners === 'function') { - ee.removeAllListeners = this._patchRemoveAllListeners( - ee, - ee.removeAllListeners - ); - } - return target; - } - - /** - * Patch methods that remove a given listener so that we match the "patched" - * version of that listener (the one that propagate context). - * @param ee EventEmitter instance - * @param original reference to the patched method - */ - private _patchRemoveListener(ee: PatchedEventEmitter, original: Function) { - return function (this: {}, event: string, listener: Func) { - if ( - ee.__ot_listeners === undefined || - ee.__ot_listeners[event] === undefined - ) { - return original.call(this, event, listener); - } - const events = ee.__ot_listeners[event]; - const patchedListener = events.get(listener); - return original.call(this, event, patchedListener || listener); - }; - } - - /** - * Patch methods that remove all listeners so we remove our - * internal references for a given event. - * @param ee EventEmitter instance - * @param original reference to the patched method - */ - private _patchRemoveAllListeners( - ee: PatchedEventEmitter, - original: Function - ) { - return function (this: {}, event: string) { - if ( - ee.__ot_listeners === undefined || - ee.__ot_listeners[event] === undefined - ) { - return original.call(this, event); - } - delete ee.__ot_listeners[event]; - return original.call(this, event); - }; - } - - /** - * Patch methods on an event emitter instance that can add listeners so we - * can force them to propagate a given context. - * @param ee EventEmitter instance - * @param original reference to the patched method - * @param [context] context to propagate when calling listeners - */ - private _patchAddListener( - ee: PatchedEventEmitter, - original: Function, - context: Context - ) { - const contextManager = this; - return function (this: {}, event: string, listener: Func) { - if (ee.__ot_listeners === undefined) ee.__ot_listeners = {}; - let listeners = ee.__ot_listeners[event]; - if (listeners === undefined) { - listeners = new WeakMap(); - ee.__ot_listeners[event] = listeners; - } - const patchedListener = contextManager.bind(listener, context); - // store a weak reference of the user listener to ours - listeners.set(listener, patchedListener); - return original.call(this, event, patchedListener); - }; - } - /** * Init hook will be called when userland create a async context, setting the * context as the current one if it exist. diff --git a/packages/opentelemetry-context-async-hooks/src/AsyncLocalStorageContextManager.ts b/packages/opentelemetry-context-async-hooks/src/AsyncLocalStorageContextManager.ts new file mode 100644 index 0000000000..c737dd834c --- /dev/null +++ b/packages/opentelemetry-context-async-hooks/src/AsyncLocalStorageContextManager.ts @@ -0,0 +1,48 @@ +/* + * Copyright The OpenTelemetry Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Context } from '@opentelemetry/context-base'; +import { AsyncLocalStorage } from 'async_hooks'; +import { BaseContextManager } from './BaseContextManager'; + +export class AsyncLocalStorageContextManager extends BaseContextManager { + private _asyncLocalStorage: AsyncLocalStorage; + + constructor() { + super(); + this._asyncLocalStorage = new AsyncLocalStorage(); + } + + active(): Context { + return this._asyncLocalStorage.getStore() ?? Context.ROOT_CONTEXT; + } + + with ReturnType>( + context: Context, + fn: T + ): ReturnType { + return this._asyncLocalStorage.run(context, fn) as ReturnType; + } + + enable(): this { + return this; + } + + disable(): this { + this._asyncLocalStorage.disable(); + return this; + } +} diff --git a/packages/opentelemetry-context-async-hooks/src/BaseContextManager.ts b/packages/opentelemetry-context-async-hooks/src/BaseContextManager.ts new file mode 100644 index 0000000000..dda3b659c6 --- /dev/null +++ b/packages/opentelemetry-context-async-hooks/src/BaseContextManager.ts @@ -0,0 +1,188 @@ +/* + * Copyright The OpenTelemetry Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { ContextManager, Context } from '@opentelemetry/context-base'; +import { EventEmitter } from 'events'; + +type Func = (...args: unknown[]) => T; + +type PatchedEventEmitter = { + /** + * Store a map for each event of all original listener and their "patched" + * version so when the listener is removed by the user, we remove the + * corresponding "patched" function. + */ + __ot_listeners?: { [name: string]: WeakMap, Func> }; +} & EventEmitter; + +const ADD_LISTENER_METHODS = [ + 'addListener' as 'addListener', + 'on' as 'on', + 'once' as 'once', + 'prependListener' as 'prependListener', + 'prependOnceListener' as 'prependOnceListener', +]; + +export abstract class BaseContextManager implements ContextManager { + abstract active(): Context; + + abstract with ReturnType>( + context: Context, + fn: T + ): ReturnType; + + bind(target: T, context?: Context): T { + // if no specific context to propagate is given, we use the current one + if (context === undefined) { + context = this.active(); + } + if (target instanceof EventEmitter) { + return this._bindEventEmitter(target, context); + } else if (typeof target === 'function') { + return this._bindFunction(target, context); + } + return target; + } + + abstract enable(): this; + + abstract disable(): this; + + private _bindFunction(target: T, context: Context): T { + const manager = this; + const contextWrapper = function (this: {}, ...args: unknown[]) { + return manager.with(context, () => target.apply(this, args)); + }; + Object.defineProperty(contextWrapper, 'length', { + enumerable: false, + configurable: true, + writable: false, + value: target.length, + }); + /** + * It isn't possible to tell Typescript that contextWrapper is the same as T + * so we forced to cast as any here. + */ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + return contextWrapper as any; + } + + /** + * By default, EventEmitter call their callback with their context, which we do + * not want, instead we will bind a specific context to all callbacks that + * go through it. + * @param target EventEmitter a instance of EventEmitter to patch + * @param context the context we want to bind + */ + private _bindEventEmitter( + target: T, + context: Context + ): T { + const ee = (target as unknown) as PatchedEventEmitter; + if (ee.__ot_listeners !== undefined) return target; + ee.__ot_listeners = {}; + + // patch methods that add a listener to propagate context + ADD_LISTENER_METHODS.forEach(methodName => { + if (ee[methodName] === undefined) return; + ee[methodName] = this._patchAddListener(ee, ee[methodName], context); + }); + // patch methods that remove a listener + if (typeof ee.removeListener === 'function') { + ee.removeListener = this._patchRemoveListener(ee, ee.removeListener); + } + if (typeof ee.off === 'function') { + ee.off = this._patchRemoveListener(ee, ee.off); + } + // patch method that remove all listeners + if (typeof ee.removeAllListeners === 'function') { + ee.removeAllListeners = this._patchRemoveAllListeners( + ee, + ee.removeAllListeners + ); + } + return target; + } + + /** + * Patch methods that remove a given listener so that we match the "patched" + * version of that listener (the one that propagate context). + * @param ee EventEmitter instance + * @param original reference to the patched method + */ + private _patchRemoveListener(ee: PatchedEventEmitter, original: Function) { + return function (this: {}, event: string, listener: Func) { + if ( + ee.__ot_listeners === undefined || + ee.__ot_listeners[event] === undefined + ) { + return original.call(this, event, listener); + } + const events = ee.__ot_listeners[event]; + const patchedListener = events.get(listener); + return original.call(this, event, patchedListener || listener); + }; + } + + /** + * Patch methods that remove all listeners so we remove our + * internal references for a given event. + * @param ee EventEmitter instance + * @param original reference to the patched method + */ + private _patchRemoveAllListeners( + ee: PatchedEventEmitter, + original: Function + ) { + return function (this: {}, event: string) { + if ( + ee.__ot_listeners === undefined || + ee.__ot_listeners[event] === undefined + ) { + return original.call(this, event); + } + delete ee.__ot_listeners[event]; + return original.call(this, event); + }; + } + + /** + * Patch methods on an event emitter instance that can add listeners so we + * can force them to propagate a given context. + * @param ee EventEmitter instance + * @param original reference to the patched method + * @param [context] context to propagate when calling listeners + */ + private _patchAddListener( + ee: PatchedEventEmitter, + original: Function, + context: Context + ) { + const contextManager = this; + return function (this: {}, event: string, listener: Func) { + if (ee.__ot_listeners === undefined) ee.__ot_listeners = {}; + let listeners = ee.__ot_listeners[event]; + if (listeners === undefined) { + listeners = new WeakMap(); + ee.__ot_listeners[event] = listeners; + } + const patchedListener = contextManager.bind(listener, context); + // store a weak reference of the user listener to ours + listeners.set(listener, patchedListener); + return original.call(this, event, patchedListener); + }; + } +} diff --git a/packages/opentelemetry-context-async-hooks/src/index.ts b/packages/opentelemetry-context-async-hooks/src/index.ts index 9213cb8ae7..94ad15e64a 100644 --- a/packages/opentelemetry-context-async-hooks/src/index.ts +++ b/packages/opentelemetry-context-async-hooks/src/index.ts @@ -15,3 +15,4 @@ */ export * from './AsyncHooksContextManager'; +export * from './AsyncLocalStorageContextManager'; diff --git a/packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts b/packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts index 15814dd196..674639081d 100644 --- a/packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts +++ b/packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts @@ -15,339 +15,366 @@ */ import * as assert from 'assert'; -import { AsyncHooksContextManager } from '../src'; +import { + AsyncHooksContextManager, + AsyncLocalStorageContextManager, +} from '../src'; import { EventEmitter } from 'events'; import { Context } from '@opentelemetry/context-base'; -describe('AsyncHooksContextManager', () => { - let contextManager: AsyncHooksContextManager; - const key1 = Context.createKey('test key 1'); - - beforeEach(() => { - contextManager = new AsyncHooksContextManager(); - contextManager.enable(); - }); - - afterEach(() => { - contextManager.disable(); - }); - - describe('.enable()', () => { - it('should work', () => { - assert.doesNotThrow(() => { - contextManager = new AsyncHooksContextManager(); - assert( - contextManager.enable() === contextManager, - 'should return this' - ); - }); +for (const contextManagerClass of [ + AsyncHooksContextManager, + AsyncLocalStorageContextManager, +]) { + describe(contextManagerClass.name, () => { + let contextManager: + | AsyncHooksContextManager + | AsyncLocalStorageContextManager; + const key1 = Context.createKey('test key 1'); + + before(function () { + if ( + contextManagerClass.name === 'AsyncLocalStorageContextManager' && + (process.version.startsWith('8.') || process.version.startsWith('10.')) + ) { + this.skip(); + } }); - }); - describe('.disable()', () => { - it('should work', () => { - assert.doesNotThrow(() => { - assert( - contextManager.disable() === contextManager, - 'should return this' - ); - }); + beforeEach(() => { + contextManager = new contextManagerClass(); contextManager.enable(); }); - }); - describe('.with()', () => { - it('should run the callback (null as target)', done => { - contextManager.with(Context.ROOT_CONTEXT, done); + afterEach(() => { + contextManager.disable(); }); - it('should run the callback (object as target)', done => { - const test = Context.ROOT_CONTEXT.setValue(key1, 1); - contextManager.with(test, () => { - assert.strictEqual( - contextManager.active(), - test, - 'should have context' - ); - return done(); + describe('.enable()', () => { + it('should work', () => { + assert.doesNotThrow(() => { + contextManager = new contextManagerClass(); + assert( + contextManager.enable() === contextManager, + 'should return this' + ); + }); }); }); - it('should run the callback (when disabled)', done => { - contextManager.disable(); - contextManager.with(Context.ROOT_CONTEXT, () => { + describe('.disable()', () => { + it('should work', () => { + assert.doesNotThrow(() => { + assert( + contextManager.disable() === contextManager, + 'should return this' + ); + }); contextManager.enable(); - return done(); }); }); - it('should rethrow errors', done => { - assert.throws(() => { + describe('.with()', () => { + it('should run the callback (null as target)', done => { + contextManager.with(Context.ROOT_CONTEXT, done); + }); + + it('should run the callback (object as target)', done => { + const test = Context.ROOT_CONTEXT.setValue(key1, 1); + contextManager.with(test, () => { + assert.strictEqual( + contextManager.active(), + test, + 'should have context' + ); + return done(); + }); + }); + + it('should run the callback (when disabled)', done => { + contextManager.disable(); contextManager.with(Context.ROOT_CONTEXT, () => { - throw new Error('This should be rethrown'); + contextManager.enable(); + return done(); }); }); - return done(); - }); - it('should finally restore an old context', done => { - const ctx1 = Context.ROOT_CONTEXT.setValue(key1, 'ctx1'); - const ctx2 = Context.ROOT_CONTEXT.setValue(key1, 'ctx2'); - contextManager.with(ctx1, () => { - assert.strictEqual(contextManager.active(), ctx1); - contextManager.with(ctx2, () => { - assert.strictEqual(contextManager.active(), ctx2); + it('should rethrow errors', done => { + assert.throws(() => { + contextManager.with(Context.ROOT_CONTEXT, () => { + throw new Error('This should be rethrown'); + }); }); - assert.strictEqual(contextManager.active(), ctx1); return done(); }); - }); - it('should finally restore an old context', done => { - const ctx1 = Context.ROOT_CONTEXT.setValue(key1, 'ctx1'); - contextManager.with(ctx1, () => { - assert.strictEqual(contextManager.active(), ctx1); - setTimeout(() => { + it('should finally restore an old context', done => { + const ctx1 = Context.ROOT_CONTEXT.setValue(key1, 'ctx1'); + const ctx2 = Context.ROOT_CONTEXT.setValue(key1, 'ctx2'); + contextManager.with(ctx1, () => { + assert.strictEqual(contextManager.active(), ctx1); + contextManager.with(ctx2, () => { + assert.strictEqual(contextManager.active(), ctx2); + }); assert.strictEqual(contextManager.active(), ctx1); return done(); }); }); - }); - it('async function called from nested "with" sync function should return nested context', done => { - const scope1 = '1' as any; - const scope2 = '2' as any; + it('should finally restore an old context', done => { + const ctx1 = Context.ROOT_CONTEXT.setValue(key1, 'ctx1'); + contextManager.with(ctx1, () => { + assert.strictEqual(contextManager.active(), ctx1); + setTimeout(() => { + assert.strictEqual(contextManager.active(), ctx1); + return done(); + }); + }); + }); + + it('async function called from nested "with" sync function should return nested context', done => { + const scope1 = '1' as any; + const scope2 = '2' as any; - const asyncFuncCalledDownstreamFromSync = async () => { - await (async () => {})(); - assert.strictEqual(contextManager.active(), scope2); - return done(); - }; + const asyncFuncCalledDownstreamFromSync = async () => { + await (async () => {})(); + assert.strictEqual(contextManager.active(), scope2); + return done(); + }; - contextManager.with(scope1, () => { - assert.strictEqual(contextManager.active(), scope1); - contextManager.with(scope2, () => asyncFuncCalledDownstreamFromSync()); - assert.strictEqual(contextManager.active(), scope1); + contextManager.with(scope1, () => { + assert.strictEqual(contextManager.active(), scope1); + contextManager.with(scope2, () => + asyncFuncCalledDownstreamFromSync() + ); + assert.strictEqual(contextManager.active(), scope1); + }); + assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); }); - assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); - }); - it('should not loose the context', done => { - const scope1 = '1' as any; + it('should not loose the context', done => { + const scope1 = '1' as any; - contextManager.with(scope1, async () => { - assert.strictEqual(contextManager.active(), scope1); - await new Promise(resolve => setTimeout(resolve, 100)); - assert.strictEqual(contextManager.active(), scope1); - return done(); + contextManager.with(scope1, async () => { + assert.strictEqual(contextManager.active(), scope1); + await new Promise(resolve => setTimeout(resolve, 100)); + assert.strictEqual(contextManager.active(), scope1); + return done(); + }); + assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); }); - assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); - }); - it('should correctly restore context using async/await', async () => { - const scope1 = '1' as any; - const scope2 = '2' as any; - const scope3 = '3' as any; - const scope4 = '4' as any; + it('should correctly restore context using async/await', async () => { + const scope1 = '1' as any; + const scope2 = '2' as any; + const scope3 = '3' as any; + const scope4 = '4' as any; - await contextManager.with(scope1, async () => { - assert.strictEqual(contextManager.active(), scope1); - await contextManager.with(scope2, async () => { - assert.strictEqual(contextManager.active(), scope2); - await contextManager.with(scope3, async () => { - assert.strictEqual(contextManager.active(), scope3); - await contextManager.with(scope4, async () => { - assert.strictEqual(contextManager.active(), scope4); + await contextManager.with(scope1, async () => { + assert.strictEqual(contextManager.active(), scope1); + await contextManager.with(scope2, async () => { + assert.strictEqual(contextManager.active(), scope2); + await contextManager.with(scope3, async () => { + assert.strictEqual(contextManager.active(), scope3); + await contextManager.with(scope4, async () => { + assert.strictEqual(contextManager.active(), scope4); + }); + assert.strictEqual(contextManager.active(), scope3); }); - assert.strictEqual(contextManager.active(), scope3); + assert.strictEqual(contextManager.active(), scope2); }); - assert.strictEqual(contextManager.active(), scope2); + assert.strictEqual(contextManager.active(), scope1); }); - assert.strictEqual(contextManager.active(), scope1); + assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); }); - assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); - }); - it('should works with multiple concurrent operations', done => { - const scope1 = '1' as any; - const scope2 = '2' as any; - const scope3 = '3' as any; - const scope4 = '4' as any; - let scope4Called = false; - - contextManager.with(scope1, async () => { - assert.strictEqual(contextManager.active(), scope1); - setTimeout(async () => { - await contextManager.with(scope3, async () => { - assert.strictEqual(contextManager.active(), scope3); - }); + it('should works with multiple concurrent operations', done => { + const scope1 = '1' as any; + const scope2 = '2' as any; + const scope3 = '3' as any; + const scope4 = '4' as any; + let scope4Called = false; + + contextManager.with(scope1, async () => { assert.strictEqual(contextManager.active(), scope1); - assert.strictEqual(scope4Called, true); - return done(); - }, 100); - assert.strictEqual(contextManager.active(), scope1); - }); - assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); - contextManager.with(scope2, async () => { - assert.strictEqual(contextManager.active(), scope2); - setTimeout(() => { - contextManager.with(scope4, async () => { - assert.strictEqual(contextManager.active(), scope4); - scope4Called = true; - }); + setTimeout(async () => { + await contextManager.with(scope3, async () => { + assert.strictEqual(contextManager.active(), scope3); + }); + assert.strictEqual(contextManager.active(), scope1); + assert.strictEqual(scope4Called, true); + return done(); + }, 100); + assert.strictEqual(contextManager.active(), scope1); + }); + assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); + contextManager.with(scope2, async () => { + assert.strictEqual(contextManager.active(), scope2); + setTimeout(() => { + contextManager.with(scope4, async () => { + assert.strictEqual(contextManager.active(), scope4); + scope4Called = true; + }); + assert.strictEqual(contextManager.active(), scope2); + }, 20); assert.strictEqual(contextManager.active(), scope2); - }, 20); - assert.strictEqual(contextManager.active(), scope2); + }); + assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); }); - assert.strictEqual(contextManager.active(), Context.ROOT_CONTEXT); - }); - }); - - describe('.bind(function)', () => { - it('should return the same target (when enabled)', () => { - const test = { a: 1 }; - assert.deepStrictEqual( - contextManager.bind(test, Context.ROOT_CONTEXT), - test - ); - }); - - it('should return the same target (when disabled)', () => { - contextManager.disable(); - const test = { a: 1 }; - assert.deepStrictEqual( - contextManager.bind(test, Context.ROOT_CONTEXT), - test - ); - contextManager.enable(); }); - it('should return current context (when enabled)', done => { - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const fn = contextManager.bind(() => { - assert.strictEqual( - contextManager.active(), - context, - 'should have context' + describe('.bind(function)', () => { + it('should return the same target (when enabled)', () => { + const test = { a: 1 }; + assert.deepStrictEqual( + contextManager.bind(test, Context.ROOT_CONTEXT), + test ); - return done(); - }, context); - fn(); - }); + }); - /** - * Even if asynchooks is disabled, the context propagation will - * still works but it might be lost after any async op. - */ - it('should return current context (when disabled)', done => { - contextManager.disable(); - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const fn = contextManager.bind(() => { - assert.strictEqual( - contextManager.active(), - context, - 'should have context' + it('should return the same target (when disabled)', () => { + contextManager.disable(); + const test = { a: 1 }; + assert.deepStrictEqual( + contextManager.bind(test, Context.ROOT_CONTEXT), + test ); - return done(); - }, context); - fn(); - }); + contextManager.enable(); + }); - it('should fail to return current context with async op', done => { - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const fn = contextManager.bind(() => { - assert.strictEqual(contextManager.active(), context); - setTimeout(() => { + it('should return current context (when enabled)', done => { + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const fn = contextManager.bind(() => { assert.strictEqual( contextManager.active(), context, - 'should have no context' + 'should have context' ); return done(); - }, 100); - }, context); - fn(); - }); - }); + }, context); + fn(); + }); - describe('.bind(event-emitter)', () => { - it('should return the same target (when enabled)', () => { - const ee = new EventEmitter(); - assert.deepStrictEqual(contextManager.bind(ee, Context.ROOT_CONTEXT), ee); - }); + /** + * Even if asynchooks is disabled, the context propagation will + * still works but it might be lost after any async op. + */ + it('should return current context (when disabled)', done => { + contextManager.disable(); + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const fn = contextManager.bind(() => { + assert.strictEqual( + contextManager.active(), + context, + 'should have context' + ); + return done(); + }, context); + fn(); + }); - it('should return the same target (when disabled)', () => { - const ee = new EventEmitter(); - contextManager.disable(); - assert.deepStrictEqual(contextManager.bind(ee, Context.ROOT_CONTEXT), ee); + it('should fail to return current context with async op', done => { + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const fn = contextManager.bind(() => { + assert.strictEqual(contextManager.active(), context); + setTimeout(() => { + assert.strictEqual( + contextManager.active(), + context, + 'should have no context' + ); + return done(); + }, 100); + }, context); + fn(); + }); }); - it('should return current context and removeListener (when enabled)', done => { - const ee = new EventEmitter(); - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const patchedEe = contextManager.bind(ee, context); - const handler = () => { - assert.deepStrictEqual(contextManager.active(), context); - patchedEe.removeListener('test', handler); - assert.strictEqual(patchedEe.listeners('test').length, 0); - return done(); - }; - patchedEe.on('test', handler); - assert.strictEqual(patchedEe.listeners('test').length, 1); - patchedEe.emit('test'); - }); + describe('.bind(event-emitter)', () => { + it('should return the same target (when enabled)', () => { + const ee = new EventEmitter(); + assert.deepStrictEqual( + contextManager.bind(ee, Context.ROOT_CONTEXT), + ee + ); + }); - it('should return current context and removeAllListener (when enabled)', done => { - const ee = new EventEmitter(); - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const patchedEe = contextManager.bind(ee, context); - const handler = () => { - assert.deepStrictEqual(contextManager.active(), context); - patchedEe.removeAllListeners('test'); - assert.strictEqual(patchedEe.listeners('test').length, 0); - return done(); - }; - patchedEe.on('test', handler); - assert.strictEqual(patchedEe.listeners('test').length, 1); - patchedEe.emit('test'); - }); + it('should return the same target (when disabled)', () => { + const ee = new EventEmitter(); + contextManager.disable(); + assert.deepStrictEqual( + contextManager.bind(ee, Context.ROOT_CONTEXT), + ee + ); + }); - /** - * Even if asynchooks is disabled, the context propagation will - * still works but it might be lost after any async op. - */ - it('should return context (when disabled)', done => { - contextManager.disable(); - const ee = new EventEmitter(); - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const patchedEe = contextManager.bind(ee, context); - const handler = () => { - assert.deepStrictEqual(contextManager.active(), context); - patchedEe.removeListener('test', handler); - assert.strictEqual(patchedEe.listeners('test').length, 0); - return done(); - }; - patchedEe.on('test', handler); - assert.strictEqual(patchedEe.listeners('test').length, 1); - patchedEe.emit('test'); - }); + it('should return current context and removeListener (when enabled)', done => { + const ee = new EventEmitter(); + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const patchedEe = contextManager.bind(ee, context); + const handler = () => { + assert.deepStrictEqual(contextManager.active(), context); + patchedEe.removeListener('test', handler); + assert.strictEqual(patchedEe.listeners('test').length, 0); + return done(); + }; + patchedEe.on('test', handler); + assert.strictEqual(patchedEe.listeners('test').length, 1); + patchedEe.emit('test'); + }); - it('should not return current context with async op', done => { - const ee = new EventEmitter(); - const context = Context.ROOT_CONTEXT.setValue(key1, 1); - const patchedEe = contextManager.bind(ee, context); - const handler = () => { - assert.deepStrictEqual(contextManager.active(), context); - setImmediate(() => { + it('should return current context and removeAllListener (when enabled)', done => { + const ee = new EventEmitter(); + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const patchedEe = contextManager.bind(ee, context); + const handler = () => { assert.deepStrictEqual(contextManager.active(), context); patchedEe.removeAllListeners('test'); assert.strictEqual(patchedEe.listeners('test').length, 0); return done(); - }); - }; - patchedEe.on('test', handler); - assert.strictEqual(patchedEe.listeners('test').length, 1); - patchedEe.emit('test'); + }; + patchedEe.on('test', handler); + assert.strictEqual(patchedEe.listeners('test').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. + */ + it('should return context (when disabled)', done => { + contextManager.disable(); + const ee = new EventEmitter(); + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const patchedEe = contextManager.bind(ee, context); + const handler = () => { + assert.deepStrictEqual(contextManager.active(), context); + patchedEe.removeListener('test', handler); + assert.strictEqual(patchedEe.listeners('test').length, 0); + return done(); + }; + patchedEe.on('test', handler); + assert.strictEqual(patchedEe.listeners('test').length, 1); + patchedEe.emit('test'); + }); + + it('should not return current context with async op', done => { + const ee = new EventEmitter(); + const context = Context.ROOT_CONTEXT.setValue(key1, 1); + const patchedEe = contextManager.bind(ee, context); + const handler = () => { + assert.deepStrictEqual(contextManager.active(), context); + setImmediate(() => { + assert.deepStrictEqual(contextManager.active(), context); + patchedEe.removeAllListeners('test'); + assert.strictEqual(patchedEe.listeners('test').length, 0); + return done(); + }); + }; + patchedEe.on('test', handler); + assert.strictEqual(patchedEe.listeners('test').length, 1); + patchedEe.emit('test'); + }); }); }); -}); +}