From b3f41f6abb84795786c0fa502959a55fa23c644f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Johannes=20W=C3=BCrbach?= <johannes.wuerbach@googlemail.com>
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           | 81 ++-----------------
 1 file changed, 6 insertions(+), 75 deletions(-)

diff --git a/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts b/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts
index 0d22c1fd5ee..23087ba4bbf 100644
--- a/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts
+++ b/packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts
@@ -15,7 +15,7 @@
  */
 
 import { ContextManager, Context } from '@opentelemetry/context-base';
-import * as asyncHooks from 'async_hooks';
+import { AsyncLocalStorage } from 'async_hooks';
 import { EventEmitter } from 'events';
 
 type Func<T> = (...args: unknown[]) => T;
@@ -38,34 +38,21 @@ const ADD_LISTENER_METHODS = [
 ];
 
 export class AsyncHooksContextManager implements ContextManager {
-  private _asyncHook: asyncHooks.AsyncHook;
-  private _contexts: Map<number, Context | undefined> = new Map();
-  private _stack: Array<Context | undefined> = [];
+  private _asyncLocalStorage: AsyncLocalStorage<Context>;
 
   constructor() {
-    this._asyncHook = asyncHooks.createHook({
-      init: this._init.bind(this),
-      before: this._before.bind(this),
-      after: this._after.bind(this),
-      destroy: this._destroy.bind(this),
-      promiseResolve: this._destroy.bind(this),
-    });
+    this._asyncLocalStorage = new AsyncLocalStorage();
   }
 
   active(): Context {
-    return this._stack[this._stack.length - 1] ?? Context.ROOT_CONTEXT;
+    return this._asyncLocalStorage.getStore() ?? Context.ROOT_CONTEXT;
   }
 
   with<T extends (...args: unknown[]) => ReturnType<T>>(
     context: Context,
     fn: T
   ): ReturnType<T> {
-    this._enterContext(context);
-    try {
-      return fn();
-    } finally {
-      this._exitContext();
-    }
+    return this._asyncLocalStorage.run(context, fn) as ReturnType<T>;
   }
 
   bind<T>(target: T, context?: Context): T {
@@ -82,14 +69,11 @@ export class AsyncHooksContextManager implements ContextManager {
   }
 
   enable(): this {
-    this._asyncHook.enable();
     return this;
   }
 
   disable(): this {
-    this._asyncHook.disable();
-    this._contexts.clear();
-    this._stack = [];
+    this._asyncLocalStorage.disable();
     return this;
   }
 
@@ -217,57 +201,4 @@ export class AsyncHooksContextManager implements ContextManager {
       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.
-   * @param uid id of the async context
-   */
-  private _init(uid: number) {
-    const context = this._stack[this._stack.length - 1];
-    if (context !== undefined) {
-      this._contexts.set(uid, context);
-    }
-  }
-
-  /**
-   * Destroy hook will be called when a given context is no longer used so we can
-   * remove its attached context.
-   * @param uid uid of the async context
-   */
-  private _destroy(uid: number) {
-    this._contexts.delete(uid);
-  }
-
-  /**
-   * Before hook is called just beforing executing a async context.
-   * @param uid uid of the async context
-   */
-  private _before(uid: number) {
-    const context = this._contexts.get(uid);
-    if (context !== undefined) {
-      this._enterContext(context);
-    }
-  }
-
-  /**
-   * After hook is called just after completing the execution of a async context.
-   */
-  private _after() {
-    this._exitContext();
-  }
-
-  /**
-   * Set the given context as active
-   */
-  private _enterContext(context: Context) {
-    this._stack.push(context);
-  }
-
-  /**
-   * Remove the context at the root of the stack
-   */
-  private _exitContext() {
-    this._stack.pop();
-  }
 }