Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 5ce2233

Browse files
bergundybergundy
committed
feat: Add Workflow AsyncLocalStorage
1 parent 72bdb8f commit 5ce2233

File tree

6 files changed

+186
-55
lines changed

6 files changed

+186
-55
lines changed

‎packages/test/src/test-workflows.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1182,6 +1182,8 @@ test('cancellation-error-is-propagated', async (t) => {
11821182
CancelledError: Cancelled
11831183
at CancellationScope.cancel
11841184
at eval
1185+
at eval
1186+
at AsyncLocalStorage.run
11851187
at CancellationScope.run
11861188
at Function.cancellable
11871189
at Object.main

‎packages/workflow/src/async-local-storage.ts

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
import { HookManager, HookType } from './promise-hooks';
2+
3+
/**
4+
* Workflow isolate (mostly compatible) port of Node's {@link https://nodejs.org/docs/latest/api/async_hooks.html#async_hooks_class_asynclocalstorage | AsyncLocalStorage}.
5+
*
6+
* **IMPORTANT** Works for Promises (and async functions) only, not for callbacks.
7+
*
8+
* @example
9+
* ```ts
10+
* const storage = new AsyncLocalStorage<number>();
11+
*
12+
* storage.run(1, async () => {
13+
* storage.getStore(); // 1
14+
* await sleep(20);
15+
* storage.getStore(); // 1
16+
* setTimeout(() => {
17+
* storage.getStore(); // undefined
18+
* }, 20);
19+
* });
20+
* ```
21+
*/
22+
export class AsyncLocalStorage<T> {
23+
/** Used by AsyncLocalStorage.run and PromiseHooks to track the current store */
24+
protected readonly stack = new Array<T>();
25+
26+
constructor() {
27+
HookManager.instance.register(this.hook);
28+
}
29+
30+
/**
31+
* Disables the instance of AsyncLocalStorage. Subsequent calls to `asyncLocalStorage.getStore()` **might** return undefined.
32+
*
33+
* Calling `asyncLocalStorage.disable()` is required before the asyncLocalStorage can be garbage collected.
34+
* This does not apply to stores provided by the asyncLocalStorage, as those objects are garbage collected
35+
* along with the corresponding async resources.
36+
*
37+
* Use this method when the asyncLocalStorage is not in use anymore in the current process.
38+
*/
39+
public disable(): void {
40+
HookManager.instance.deregister(this.hook);
41+
}
42+
43+
/**
44+
* Returns the current store.
45+
*
46+
* If called outside of an asynchronous context initialized by calling `asyncLocalStorage.run()` it returns undefined.
47+
*/
48+
getStore(): T | undefined {
49+
if (this.stack.length === 0) {
50+
return undefined;
51+
}
52+
return this.stack[this.stack.length - 1];
53+
}
54+
55+
/**
56+
* Runs a function synchronously within a context and returns its return value.
57+
* The store is not accessible outside of the callback function. The store is accessible to any
58+
* asynchronous operations created within the callback.
59+
*
60+
* The optional args are passed to the callback function.
61+
*/
62+
public run<A extends any[], R>(store: T, callback: (...args: A) => R, ...args: A): R {
63+
this.stack.push(store);
64+
try {
65+
return callback(...args);
66+
} finally {
67+
this.stack.pop();
68+
}
69+
}
70+
71+
/**
72+
* Hook into the v8 PromiseHook callback to keep track of the current store.
73+
*
74+
* This function is bound to the class instance so it can be used as a unique key for
75+
* registration and promise data mapping.
76+
*/
77+
protected hook = (t: HookType, p: Promise<any>): void => {
78+
switch (t) {
79+
// When a Promise is created associate it with a the current store.
80+
case 'init':
81+
HookManager.instance.setPromiseData(p, this.hook, this.getStore());
82+
break;
83+
// Called at the beginning of the PromiseReactionJob,
84+
// p is the promise about to execute, resume its store.
85+
case 'before':
86+
this.stack.push(HookManager.instance.getPromiseData(p, this.hook) as T);
87+
break;
88+
// Called at the end of the PromiseReactionJob,
89+
// pop the current Promise off the store stack.
90+
case 'after':
91+
this.stack.pop();
92+
break;
93+
}
94+
};
95+
}

‎packages/workflow/src/cancellation-scope.ts

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { AsyncLocalStorage } from './async-local-storage';
12
import { CancelledError, IllegalStateError } from './errors';
23

34
/** Magic symbol used to create the root scope - intentionally not exported */
@@ -121,13 +122,12 @@ export class CancellationScope {
121122
* @return the result of `fn`
122123
*/
123124
run<T>(fn: () => Promise<T>): Promise<T> {
124-
pushScope(this);
125-
if (this.timeout) {
126-
sleep(this.timeout).then(() => this.cancel());
127-
}
128-
const promise = fn();
129-
popScope();
130-
return promise;
125+
return storage.run(this, async () => {
126+
if (this.timeout) {
127+
sleep(this.timeout).then(() => this.cancel());
128+
}
129+
return await fn();
130+
});
131131
}
132132

133133
/**
@@ -141,11 +141,7 @@ export class CancellationScope {
141141
* Get the current "active" scope
142142
*/
143143
static current(): CancellationScope {
144-
const scope = scopeStack[scopeStack.length - 1];
145-
if (scope === undefined) {
146-
throw new IllegalStateError('No scopes in stack');
147-
}
148-
return scope;
144+
return storage.getStore() ?? ROOT_SCOPE;
149145
}
150146

151147
/** Alias to `new CancellationScope({ cancellable: true }).run(fn)` */
@@ -164,20 +160,10 @@ export class CancellationScope {
164160
}
165161
}
166162

167-
/** There can only be one of these */
168-
export const ROOT_SCOPE = new CancellationScope({ cancellable: true, parent: NO_PARENT as any });
169-
/** Used by CancellationScope.run and PromiseHooks to track the current scope */
170-
const scopeStack: CancellationScope[] = [ROOT_SCOPE];
171-
172-
/** Push a scope onto the scope stack */
173-
export function pushScope(scope: CancellationScope): void {
174-
scopeStack.push(scope);
175-
}
163+
const storage = new AsyncLocalStorage<CancellationScope>();
176164

177-
/** Pop the last scope off the scope stack */
178-
export function popScope(): void {
179-
scopeStack.pop();
180-
}
165+
/** There can only be one of these */
166+
export const ROOT_SCOPE = new CancellationScope({ cancellable: true, parent: NO_PARENT });
181167

182168
/** This function is here to avoid a circular dependency between this module and workflow.ts */
183169
let sleep = (_: number): Promise<void> => {

‎packages/workflow/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ export {
6666
export { CancelledError, DeterminismViolationError, IllegalStateError } from './errors';
6767
export { Context, ContextImpl, sleep, uuid4, validateActivityOptions, scheduleActivity } from './workflow';
6868
export * from './interceptors';
69+
export { AsyncLocalStorage } from './async-local-storage';
6970
export { CancellationScope, CancellationScopeOptions } from './cancellation-scope';
7071
export { Trigger } from './trigger';
7172
export { defaultDataConverter, DataConverter } from './converter/data-converter';

‎packages/workflow/src/init.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { state } from './internals';
33
import { WorkflowInterceptors } from './interceptors';
44
import { msToTs } from './time';
55
import { alea } from './alea';
6-
import { IsolateExtension, ScopeHookManager } from './promise-hooks';
6+
import { IsolateExtension, HookManager } from './promise-hooks';
77
import { DeterminismViolationError } from './errors';
88

99
export function overrideGlobals(): void {
@@ -79,5 +79,5 @@ export function initWorkflow(
7979
state.interceptors = interceptors;
8080
state.info = info;
8181
state.random = alea(randomnessSeed);
82-
new ScopeHookManager(isolateExtension);
82+
HookManager.instance.setIsolateExtension(isolateExtension);
8383
}

‎packages/workflow/src/promise-hooks.ts

Lines changed: 75 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { CancellationScope, ROOT_SCOPE, pushScope, popScope } from './cancellation-scope';
1+
import { IllegalStateError } from './errors';
22

33
/** v8 hook types */
44
export type HookType = 'init' | 'resolve' | 'before' | 'after';
@@ -10,43 +10,90 @@ export type PromiseHook = (t: HookType, p: Promise<any>, pp?: Promise<any>) => v
1010
*/
1111
export interface IsolateExtension {
1212
registerPromiseHook(hook: PromiseHook): void;
13-
/** Associate a Promise with its CancellationScope */
14-
setPromiseData(p: Promise<any>, s: CancellationScope): void;
15-
/** Get the CancellationScope associated with a Promise */
16-
getPromiseData(p: Promise<any>): CancellationScope | undefined;
13+
/** Associate a Promise with each hook's custom data */
14+
setPromiseData(p: Promise<any>, s: Map<PromiseHook, any>): void;
15+
/** Get the custom hook data associated with a Promise */
16+
getPromiseData(p: Promise<any>): Map<PromiseHook, any> | undefined;
1717
}
1818

1919
/**
2020
* Uses the v8 PromiseHook callback to track the current `CancellationScope`
2121
*/
22-
export class ScopeHookManager {
23-
protected readonly childScopes = new Map<CancellationScope, Set<CancellationScope>>();
22+
export class HookManager {
23+
protected readonly registeredHooks = new Set<PromiseHook>();
24+
/**
25+
* A reference to the native isolate extension, lazily initialized along with the Workflow
26+
*/
27+
protected isolateExtension?: IsolateExtension;
28+
29+
protected constructor() {
30+
// Prevent construction other than the singleton
31+
}
2432

25-
constructor(protected readonly isolateExtension: IsolateExtension) {
33+
// Singleton instance
34+
public static instance = new HookManager();
35+
36+
/**
37+
* To be called from the Workflow runtime library to set the native module reference
38+
*/
39+
setIsolateExtension(isolateExtension: IsolateExtension): void {
40+
this.isolateExtension = isolateExtension;
2641
isolateExtension.registerPromiseHook(this.hook.bind(this));
2742
}
2843

2944
/**
30-
* The PromiseHook implementation
31-
*
32-
* Note that the parent promise is unused as it was not found neccessary for the implementation
33-
*/
34-
hook(t: HookType, p: Promise<any>, _pp?: Promise<any>): void {
35-
switch (t) {
36-
// When a Promise is created associate it with a CancellationScope
37-
case 'init':
38-
this.isolateExtension.setPromiseData(p, CancellationScope.current());
39-
break;
40-
// Called at the beginning of the PromiseReactionJob,
41-
// p is the promise about to execute, resume its scope.
42-
case 'before':
43-
pushScope(this.isolateExtension.getPromiseData(p) || ROOT_SCOPE);
44-
break;
45-
// Called at the end of the PromiseReactionJob,
46-
// pop the current Promise off the scope stack.
47-
case 'after':
48-
popScope();
49-
break;
45+
* Helper that ensures isolateExtension has been set
46+
*/
47+
getIsolateExtension(): IsolateExtension {
48+
if (this.isolateExtension === undefined) {
49+
throw new IllegalStateError('HookManager has not been properly initialized');
50+
}
51+
return this.isolateExtension;
52+
}
53+
54+
/**
55+
* Register a single promise hook callback
56+
*/
57+
register(hook: PromiseHook): void {
58+
this.registeredHooks.add(hook);
59+
}
60+
61+
/**
62+
* Deregister a single promise hook callback
63+
*/
64+
deregister(hook: PromiseHook): void {
65+
this.registeredHooks.delete(hook);
66+
}
67+
68+
/**
69+
* The PromiseHook implementation, calls all registered hooks
70+
*/
71+
protected hook(t: HookType, p: Promise<any>, pp?: Promise<any>): void {
72+
for (const hook of this.registeredHooks) {
73+
hook(t, p, pp);
74+
}
75+
}
76+
77+
/**
78+
* Get custom promise data for a promise hook
79+
*/
80+
public getPromiseData(p: Promise<any>, hook: PromiseHook): unknown {
81+
const data = this.getIsolateExtension().getPromiseData(p);
82+
if (data) {
83+
return data.get(hook);
84+
}
85+
}
86+
87+
/**
88+
* Set custom promise data for a promise hook
89+
*/
90+
public setPromiseData(p: Promise<any>, hook: PromiseHook, data: unknown): void {
91+
const isolateExtension = this.getIsolateExtension();
92+
let mapping = isolateExtension.getPromiseData(p);
93+
if (!mapping) {
94+
mapping = new Map();
95+
isolateExtension.setPromiseData(p, mapping);
5096
}
97+
mapping.set(hook, data);
5198
}
5299
}

0 commit comments

Comments
 (0)