AsyncContext with termination

AsyncContext.ts

import InternalAsyncContextState from "./InternalAsyncContextState.js";
import InternalVariable from "./InternalVariable.js";

namespace AsyncContext {
    const state = new InternalAsyncContextState();

    export class Snapshot {
        readonly #internalSnapshot = state.currentSnapshot.clone();

        discard(): void {
            this.#internalSnapshot.discard();
        }

        run<Args extends ReadonlyArray<any>, Return>(
            fn: (...args: Args) => Return,
            ...args: Args
        ): Return {
            if (!this.#internalSnapshot.isAlive) {
                throw new TypeError(`this snapshot has already been discarded`);
            }
            return state.run(this.#internalSnapshot, fn, ...args);
        }
    }

    export type VariableInit<T> = {
        name?: string | undefined,
        defaultValue?: T | undefined,
    };

    export class Variable<T> {
        readonly #internalVariable: InternalVariable<T>;
        readonly #name: string;

        constructor({ name = "", defaultValue }: VariableInit<T> = {}) {
            this.#name = name;
            this.#internalVariable = new InternalVariable(state, { name, defaultValue });
        }

        get name(): string {
            return this.#name;
        }

        get(): T | undefined {
            return this.#internalVariable.get();
        }

        run<Args extends ReadonlyArray<any>, Return>(
            value: T,
            fn: (...args: Args) => Return,
            ...args: Args
        ): Return {
            return this.#internalVariable.run(value, fn, ...args);
        }

        runWithTermination<Args extends ReadonlyArray<any>, Return>(
            value: T,
            terminationHandler: () => void,
            fn: (...args: Args) => Return,
            ...args: Args
        ): Return {
            return this.#internalVariable.runWithTermination(
                value,
                terminationHandler,
                fn,
                ...args,
            );
        }
    }
}

export default AsyncContext;

InternalAsyncContextState.ts

import { createHook, executionAsyncId } from "node:async_hooks";
import InternalSnapshot from "./InternalSnapshot.js";

export default class InternalAsyncContextState {
    readonly #initialSnapshot = new InternalSnapshot(this, new Map());

    readonly #resourceMapping = new Map<number, InternalSnapshot>();

    readonly #hook = createHook({
        destroy: (asyncId) => {
            const snapshot = this.#resourceMapping.get(asyncId);
            snapshot?.discard();
        },
        init: (asyncId, type, triggerAsyncId, resource) => {
            const previousSnapshot
                = this.#resourceMapping.get(triggerAsyncId) ?? this.#initialSnapshot;
            this.#resourceMapping.set(asyncId, previousSnapshot.clone());
        },
    });

    constructor() {
        this.#hook.enable();
    }

    disable(): void {
        this.#hook.disable();
    }

    get currentSnapshot(): InternalSnapshot {
        return this.#resourceMapping.get(executionAsyncId()) ?? this.#initialSnapshot;
    }

    run<Args extends ReadonlyArray<any>, Return>(
        snapshot: InternalSnapshot,
        fn: (...args: Args) => Return,
        ...args: Args
    ): Return {
        const asyncId = executionAsyncId();
        const previousSnapshot
            = this.#resourceMapping.get(asyncId) ?? this.#initialSnapshot.clone();
        this.#resourceMapping.set(asyncId, snapshot);
        try {
            return fn(...args);
        } finally {
            this.#resourceMapping.set(asyncId, previousSnapshot);
        }
    }
}

InternalSnapshot.ts

import type InternalAsyncContextState from "./InternalAsyncContextState.js";
import type InternalVariable from "./InternalVariable.js";

type VariableStateInit = {
    value: any,
    terminationHandler: () => void,
};

class InternalVariableState {
    #value: {
        value: any,
        terminationHandler: () => void,
    } | null;
    #refCount: number = 0;

    constructor({ value, terminationHandler }: VariableStateInit) {
        this.#value = { value, terminationHandler };
    }

    [Symbol.for("nodejs.util.inspect.custom")](
        depth: number,
        options: import("node:util").InspectOptions,
        inspect: typeof import("node:util").inspect,
    ) {
        return inspect(
            {
                value: this.#value?.value,
                terminationHandler: this.#value?.terminationHandler,
                refCount: this.#refCount,
            },
            { depth, ...options },
        );
    }

    get value(): any {
        if (this.#value === null) {
            throw new RangeError(`value has been discarded`);
        }
        return this.#value?.value;
    }

    ref(): void {
        this.#refCount += 1;
    }

    deref(): void {
        this.#refCount -= 1;
        if (this.#refCount === 0) {
            if (this.#value === null) {
                throw new RangeError(`this should be impossible`);
            }
            const { terminationHandler } = this.#value;
            this.#value = null;
            terminationHandler();
        }
    }
}

const finalizer = new FinalizationRegistry<Map<InternalVariable<any>, InternalVariableState>>(
    (mapping) => {
        for (const variableState of mapping.values()) {
            variableState.deref();
        }
    },
);

export default class InternalSnapshot {
    readonly #state: InternalAsyncContextState;

    #maybeMapping: Map<InternalVariable<any>, InternalVariableState> | null;

    constructor(
        state: InternalAsyncContextState,
        mapping: Map<InternalVariable<any>, InternalVariableState>,
    ) {
        this.#state = state;
        this.#maybeMapping = mapping;
        for (const variableState of mapping.values()) {
            variableState.ref();
        }
        finalizer.register(this, mapping, this);
    }

    [Symbol.for("nodejs.util.inspect.custom")](
        depth: number,
        options: import("node:util").InspectOptions,
        inspect: typeof import("node:util").inspect,
    ) {
        return `InternalSnapshot ${inspect(
            {
                "#state": this.#state,
                "#mapping": this.#maybeMapping,
            },
            { depth, ...options },
        )}`;
    }

    get #mapping(): Map<InternalVariable<any>, InternalVariableState> {
        if (this.#maybeMapping === null) {
            throw new Error(`can't use InternalSnapshot after it has been discarded`);
        }
        return this.#maybeMapping;
    }

    get isAlive(): boolean {
        return this.#maybeMapping !== null;
    }

    discard(): void {
        finalizer.unregister(this);
        if (this.#maybeMapping) {
            for (const variableState of this.#mapping.values()) {
                variableState.deref();
            }
        }
        this.#maybeMapping = null;
    }

    get(variable: InternalVariable<any>): { value: any } | null {
        const state = this.#mapping.get(variable);
        if (state) {
            return { value: state.value };
        }
        return null;
    }

    clone(): InternalSnapshot {
        return new InternalSnapshot(this.#state, new Map([...this.#mapping]));
    }

    cloneWith(
        variable: InternalVariable<any>,
        value: any,
        terminationHandler: () => void,
    ): InternalSnapshot {
        const variableState = new InternalVariableState({ value, terminationHandler });
        return new InternalSnapshot(
            this.#state,
            new Map([...this.#mapping, [variable, variableState]]),
        );
    }

    run<Args extends ReadonlyArray<any>, Return>(
        fn: (...args: Args) => Return,
        ...args: Args
    ): Return {
        return this.#state.run(this, fn, ...args);
    }
}

InternalVariable.ts

import type AsyncContextState from "./InternalAsyncContextState.js";

export type AsyncVariableOptions<T> = {
    name?: string | undefined,
    defaultValue?: T | undefined,
};

export default class InternalVariable<T> {
    readonly #state: AsyncContextState;
    readonly #name: string;
    readonly #defaultValue: T | undefined;

    constructor(
        state: AsyncContextState,
        { name = "", defaultValue }: AsyncVariableOptions<T> = {},
    ) {
        this.#state = state;
        this.#name = name;
        this.#defaultValue = defaultValue;
    }

    get name(): string {
        return this.#name;
    }

    get(): T | undefined {
        const v = this.#state.currentSnapshot.get(this);
        if (v === null) {
            return this.#defaultValue;
        }
        return v.value;
    }

    #run<Args extends ReadonlyArray<any>, Return>(
        value: T,
        terminationCb: () => void,
        fn: (...args: Args) => Return,
        ...args: Args
    ): Return {
        const modifiedSnapshot = this.#state.currentSnapshot.cloneWith(this, value, terminationCb);
        try {
            return this.#state.run(modifiedSnapshot, fn, ...args);
        } finally {
            modifiedSnapshot.discard();
        }
    }

    run<Args extends ReadonlyArray<any>, Return>(
        value: T,
        fn: (...args: Args) => Return,
        ...args: Args
    ): Return {
        return this.#run(value, () => {}, fn, ...args);
    }

    runWithTermination<Args extends ReadonlyArray<any>, Return>(
        value: T,
        terminationCb: () => void,
        fn: (...args: Args) => Return,
        ...args: Args
    ): Return {
        return this.#run(value, terminationCb, fn, ...args);
    }
}