Skip to content

Commit d9e91d0

Browse files
Tom KlaverTommos0kamilkisiela
authored
Use AsyncLocalStorage for execution context if available (#2395)
Co-authored-by: Tom Klaver <[email protected]> Co-authored-by: Kamil Kisiela <[email protected]>
1 parent b6bcac0 commit d9e91d0

6 files changed

+133
-79
lines changed

‎.changeset/hip-elephants-warn.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"graphql-modules": minor
3+
---
4+
5+
Use AsyncLocalStorage for execution context if available

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

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import { AsyncLocalStorage } from 'async_hooks';
2+
import { type ExecutionContextPicker } from './execution-context.interface';
3+
4+
const executionContextStore = AsyncLocalStorage
5+
? new AsyncLocalStorage<ExecutionContextPicker>()
6+
: undefined;
7+
8+
export const executionContext: {
9+
create(picker: ExecutionContextPicker): () => void;
10+
getModuleContext: ExecutionContextPicker['getModuleContext'];
11+
getApplicationContext: ExecutionContextPicker['getApplicationContext'];
12+
} = {
13+
create(picker) {
14+
executionContextStore!.enterWith(picker);
15+
return function destroyContext() {};
16+
},
17+
getModuleContext(moduleId) {
18+
return executionContextStore!.getStore()!.getModuleContext(moduleId);
19+
},
20+
getApplicationContext() {
21+
return executionContextStore!.getStore()!.getApplicationContext();
22+
},
23+
};
24+
25+
export function enableExecutionContext() {}
26+
27+
export function getExecutionContextStore() {
28+
return executionContextStore;
29+
}

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

+79
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
import { createHook, executionAsyncId } from 'async_hooks';
2+
import { type ExecutionContextPicker } from './execution-context.interface';
3+
4+
const executionContextStore = new Map<number, ExecutionContextPicker>();
5+
const executionContextDependencyStore = new Map<number, Set<number>>();
6+
7+
const executionContextHook = createHook({
8+
init(asyncId, _, triggerAsyncId) {
9+
// Store same context data for child async resources
10+
const ctx = executionContextStore.get(triggerAsyncId);
11+
if (ctx) {
12+
const dependencies =
13+
executionContextDependencyStore.get(triggerAsyncId) ??
14+
executionContextDependencyStore
15+
.set(triggerAsyncId, new Set())
16+
.get(triggerAsyncId)!;
17+
dependencies.add(asyncId);
18+
executionContextStore.set(asyncId, ctx);
19+
}
20+
},
21+
destroy(asyncId) {
22+
if (executionContextStore.has(asyncId)) {
23+
executionContextStore.delete(asyncId);
24+
}
25+
},
26+
});
27+
28+
function destroyContextAndItsChildren(id: number) {
29+
if (executionContextStore.has(id)) {
30+
executionContextStore.delete(id);
31+
}
32+
33+
const deps = executionContextDependencyStore.get(id);
34+
35+
if (deps) {
36+
for (const dep of deps) {
37+
destroyContextAndItsChildren(dep);
38+
}
39+
executionContextDependencyStore.delete(id);
40+
}
41+
}
42+
43+
export const executionContext: {
44+
create(picker: ExecutionContextPicker): () => void;
45+
getModuleContext: ExecutionContextPicker['getModuleContext'];
46+
getApplicationContext: ExecutionContextPicker['getApplicationContext'];
47+
} = {
48+
create(picker) {
49+
const id = executionAsyncId();
50+
executionContextStore.set(id, picker);
51+
return function destroyContext() {
52+
destroyContextAndItsChildren(id);
53+
};
54+
},
55+
getModuleContext(moduleId) {
56+
const picker = executionContextStore.get(executionAsyncId())!;
57+
return picker.getModuleContext(moduleId);
58+
},
59+
getApplicationContext() {
60+
const picker = executionContextStore.get(executionAsyncId())!;
61+
return picker.getApplicationContext();
62+
},
63+
};
64+
65+
let executionContextEnabled = false;
66+
67+
export function enableExecutionContext() {
68+
if (!executionContextEnabled) {
69+
executionContextHook.enable();
70+
}
71+
}
72+
73+
export function getExecutionContextStore() {
74+
return executionContextStore;
75+
}
76+
77+
export function getExecutionContextDependencyStore() {
78+
return executionContextDependencyStore;
79+
}

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

+4
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
export interface ExecutionContextPicker {
2+
getModuleContext(moduleId: string): GraphQLModules.ModuleContext;
3+
getApplicationContext(): GraphQLModules.AppContext;
4+
}

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

+14-78
Original file line numberDiff line numberDiff line change
@@ -1,83 +1,19 @@
1-
import { createHook, executionAsyncId } from 'async_hooks';
1+
import { AsyncLocalStorage } from 'async_hooks';
22

3-
export interface ExecutionContextPicker {
4-
getModuleContext(moduleId: string): GraphQLModules.ModuleContext;
5-
getApplicationContext(): GraphQLModules.AppContext;
6-
}
3+
/*
4+
Use AsyncLocalStorage if available (available sync Node 14).
5+
Otherwise, fall back to using async_hooks.createHook
6+
*/
77

8-
const executionContextStore = new Map<number, ExecutionContextPicker>();
9-
const executionContextDependencyStore = new Map<number, Set<number>>();
8+
import * as Hooks from './execution-context-hooks';
9+
import * as Async from './execution-context-async-local-storage';
1010

11-
const executionContextHook = createHook({
12-
init(asyncId, _, triggerAsyncId) {
13-
// Store same context data for child async resources
14-
const ctx = executionContextStore.get(triggerAsyncId);
15-
if (ctx) {
16-
const dependencies =
17-
executionContextDependencyStore.get(triggerAsyncId) ??
18-
executionContextDependencyStore
19-
.set(triggerAsyncId, new Set())
20-
.get(triggerAsyncId)!;
21-
dependencies.add(asyncId);
22-
executionContextStore.set(asyncId, ctx);
23-
}
24-
},
25-
destroy(asyncId) {
26-
if (executionContextStore.has(asyncId)) {
27-
executionContextStore.delete(asyncId);
28-
}
29-
},
30-
});
11+
export type { ExecutionContextPicker } from './execution-context.interface';
3112

32-
function destroyContextAndItsChildren(id: number) {
33-
if (executionContextStore.has(id)) {
34-
executionContextStore.delete(id);
35-
}
13+
export const executionContext = AsyncLocalStorage
14+
? Async.executionContext
15+
: Hooks.executionContext;
3616

37-
const deps = executionContextDependencyStore.get(id);
38-
39-
if (deps) {
40-
for (const dep of deps) {
41-
destroyContextAndItsChildren(dep);
42-
}
43-
executionContextDependencyStore.delete(id);
44-
}
45-
}
46-
47-
export const executionContext: {
48-
create(picker: ExecutionContextPicker): () => void;
49-
getModuleContext: ExecutionContextPicker['getModuleContext'];
50-
getApplicationContext: ExecutionContextPicker['getApplicationContext'];
51-
} = {
52-
create(picker) {
53-
const id = executionAsyncId();
54-
executionContextStore.set(id, picker);
55-
return function destroyContext() {
56-
destroyContextAndItsChildren(id);
57-
};
58-
},
59-
getModuleContext(moduleId) {
60-
const picker = executionContextStore.get(executionAsyncId())!;
61-
return picker.getModuleContext(moduleId);
62-
},
63-
getApplicationContext() {
64-
const picker = executionContextStore.get(executionAsyncId())!;
65-
return picker.getApplicationContext();
66-
},
67-
};
68-
69-
let executionContextEnabled = false;
70-
71-
export function enableExecutionContext() {
72-
if (!executionContextEnabled) {
73-
executionContextHook.enable();
74-
}
75-
}
76-
77-
export function getExecutionContextStore() {
78-
return executionContextStore;
79-
}
80-
81-
export function getExecutionContextDependencyStore() {
82-
return executionContextDependencyStore;
83-
}
17+
export const enableExecutionContext = AsyncLocalStorage
18+
? () => undefined
19+
: Hooks.enableExecutionContext;

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

+2-1
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,11 @@ import {
99
InjectionToken,
1010
testkit,
1111
} from '../src';
12+
1213
import {
1314
getExecutionContextDependencyStore,
1415
getExecutionContextStore,
15-
} from '../src/application/execution-context';
16+
} from '../src/application/execution-context-hooks';
1617

1718
const posts = ['Foo', 'Bar'];
1819

0 commit comments

Comments
 (0)