Use AsyncLocalStorage for execution context if available (#2395)

Co-authored-by: Tom Klaver <tomklav@gmail.com>
Co-authored-by: Kamil Kisiela <kamil.kisiela@gmail.com>

.changeset/hip-elephants-warn.md

@@ -0,0 +1,5 @@
+---
+"graphql-modules": minor
+---
+
+Use AsyncLocalStorage for execution context if available

packages/graphql-modules/src/application/execution-context-async-local-storage.ts

@@ -0,0 +1,29 @@
+import { AsyncLocalStorage } from 'async_hooks';
+import { type ExecutionContextPicker } from './execution-context.interface';
+
+const executionContextStore = AsyncLocalStorage
+  ? new AsyncLocalStorage<ExecutionContextPicker>()
+  : undefined;
+
+export const executionContext: {
+  create(picker: ExecutionContextPicker): () => void;
+  getModuleContext: ExecutionContextPicker['getModuleContext'];
+  getApplicationContext: ExecutionContextPicker['getApplicationContext'];
+} = {
+  create(picker) {
+    executionContextStore!.enterWith(picker);
+    return function destroyContext() {};
+  },
+  getModuleContext(moduleId) {
+    return executionContextStore!.getStore()!.getModuleContext(moduleId);
+  },
+  getApplicationContext() {
+    return executionContextStore!.getStore()!.getApplicationContext();
+  },
+};
+
+export function enableExecutionContext() {}
+
+export function getExecutionContextStore() {
+  return executionContextStore;
+}

packages/graphql-modules/src/application/execution-context-hooks.ts

@@ -0,0 +1,79 @@
+import { createHook, executionAsyncId } from 'async_hooks';
+import { type ExecutionContextPicker } from './execution-context.interface';
+
+const executionContextStore = new Map<number, ExecutionContextPicker>();
+const executionContextDependencyStore = new Map<number, Set<number>>();
+
+const executionContextHook = createHook({
+  init(asyncId, _, triggerAsyncId) {
+    // Store same context data for child async resources
+    const ctx = executionContextStore.get(triggerAsyncId);
+    if (ctx) {
+      const dependencies =
+        executionContextDependencyStore.get(triggerAsyncId) ??
+        executionContextDependencyStore
+          .set(triggerAsyncId, new Set())
+          .get(triggerAsyncId)!;
+      dependencies.add(asyncId);
+      executionContextStore.set(asyncId, ctx);
+    }
+  },
+  destroy(asyncId) {
+    if (executionContextStore.has(asyncId)) {
+      executionContextStore.delete(asyncId);
+    }
+  },
+});
+
+function destroyContextAndItsChildren(id: number) {
+  if (executionContextStore.has(id)) {
+    executionContextStore.delete(id);
+  }
+
+  const deps = executionContextDependencyStore.get(id);
+
+  if (deps) {
+    for (const dep of deps) {
+      destroyContextAndItsChildren(dep);
+    }
+    executionContextDependencyStore.delete(id);
+  }
+}
+
+export const executionContext: {
+  create(picker: ExecutionContextPicker): () => void;
+  getModuleContext: ExecutionContextPicker['getModuleContext'];
+  getApplicationContext: ExecutionContextPicker['getApplicationContext'];
+} = {
+  create(picker) {
+    const id = executionAsyncId();
+    executionContextStore.set(id, picker);
+    return function destroyContext() {
+      destroyContextAndItsChildren(id);
+    };
+  },
+  getModuleContext(moduleId) {
+    const picker = executionContextStore.get(executionAsyncId())!;
+    return picker.getModuleContext(moduleId);
+  },
+  getApplicationContext() {
+    const picker = executionContextStore.get(executionAsyncId())!;
+    return picker.getApplicationContext();
+  },
+};
+
+let executionContextEnabled = false;
+
+export function enableExecutionContext() {
+  if (!executionContextEnabled) {
+    executionContextHook.enable();
+  }
+}
+
+export function getExecutionContextStore() {
+  return executionContextStore;
+}
+
+export function getExecutionContextDependencyStore() {
+  return executionContextDependencyStore;
+}

packages/graphql-modules/src/application/execution-context.interface.ts

@@ -0,0 +1,4 @@
+export interface ExecutionContextPicker {
+  getModuleContext(moduleId: string): GraphQLModules.ModuleContext;
+  getApplicationContext(): GraphQLModules.AppContext;
+}

packages/graphql-modules/src/application/execution-context.ts

@@ -1,83 +1,19 @@
-import { createHook, executionAsyncId } from 'async_hooks';
+import { AsyncLocalStorage } from 'async_hooks';
 
-export interface ExecutionContextPicker {
-  getModuleContext(moduleId: string): GraphQLModules.ModuleContext;
-  getApplicationContext(): GraphQLModules.AppContext;
-}
+/*
+  Use AsyncLocalStorage if available (available sync Node 14).
+  Otherwise, fall back to using async_hooks.createHook
+*/
 
-const executionContextStore = new Map<number, ExecutionContextPicker>();
-const executionContextDependencyStore = new Map<number, Set<number>>();
+import * as Hooks from './execution-context-hooks';
+import * as Async from './execution-context-async-local-storage';
 
-const executionContextHook = createHook({
-  init(asyncId, _, triggerAsyncId) {
-    // Store same context data for child async resources
-    const ctx = executionContextStore.get(triggerAsyncId);
-    if (ctx) {
-      const dependencies =
-        executionContextDependencyStore.get(triggerAsyncId) ??
-        executionContextDependencyStore
-          .set(triggerAsyncId, new Set())
-          .get(triggerAsyncId)!;
-      dependencies.add(asyncId);
-      executionContextStore.set(asyncId, ctx);
-    }
-  },
-  destroy(asyncId) {
-    if (executionContextStore.has(asyncId)) {
-      executionContextStore.delete(asyncId);
-    }
-  },
-});
+export type { ExecutionContextPicker } from './execution-context.interface';
 
-function destroyContextAndItsChildren(id: number) {
-  if (executionContextStore.has(id)) {
-    executionContextStore.delete(id);
-  }
+export const executionContext = AsyncLocalStorage
+  ? Async.executionContext
+  : Hooks.executionContext;
 
-  const deps = executionContextDependencyStore.get(id);
-
-  if (deps) {
-    for (const dep of deps) {
-      destroyContextAndItsChildren(dep);
-    }
-    executionContextDependencyStore.delete(id);
-  }
-}
-
-export const executionContext: {
-  create(picker: ExecutionContextPicker): () => void;
-  getModuleContext: ExecutionContextPicker['getModuleContext'];
-  getApplicationContext: ExecutionContextPicker['getApplicationContext'];
-} = {
-  create(picker) {
-    const id = executionAsyncId();
-    executionContextStore.set(id, picker);
-    return function destroyContext() {
-      destroyContextAndItsChildren(id);
-    };
-  },
-  getModuleContext(moduleId) {
-    const picker = executionContextStore.get(executionAsyncId())!;
-    return picker.getModuleContext(moduleId);
-  },
-  getApplicationContext() {
-    const picker = executionContextStore.get(executionAsyncId())!;
-    return picker.getApplicationContext();
-  },
-};
-
-let executionContextEnabled = false;
-
-export function enableExecutionContext() {
-  if (!executionContextEnabled) {
-    executionContextHook.enable();
-  }
-}
-
-export function getExecutionContextStore() {
-  return executionContextStore;
-}
-
-export function getExecutionContextDependencyStore() {
-  return executionContextDependencyStore;
-}
+export const enableExecutionContext = AsyncLocalStorage
+  ? () => undefined
+  : Hooks.enableExecutionContext;

packages/graphql-modules/tests/execution-context.spec.ts

@@ -9,10 +9,11 @@ import {
   InjectionToken,
   testkit,
 } from '../src';
+
 import {
   getExecutionContextDependencyStore,
   getExecutionContextStore,
-} from '../src/application/execution-context';
+} from '../src/application/execution-context-hooks';
 
 const posts = ['Foo', 'Bar'];