AsyncContext.ts

import {
  AsyncLocalStorage,
} from 'node:async_hooks'
import {
  IterableWeakSet,
} from './IterableWeakSet'

/**
 * Basic implementation of tc39 Stage 2 [Async Context](https://github.com/tc39/proposal-async-context)
 *
 * - https://github.com/tc39/proposal-async-context
 * - https://nodejs.org/api/async_context.html#class-asynclocalstorage
 *
 * refs:
 * - https://www.npmjs.com/package/@webfill/async-context
 */
export namespace AsyncContext {
  export class Variable<T> {
    #name: string
    #defaultValue?: T
    #storage = new AsyncLocalStorage<T>()

    constructor(options?: AsyncVariableOptions<T>) {
      this.#name = options?.name ?? ''
      this.#defaultValue = options?.defaultValue
      if (options?.defaultValue !== undefined) {
        /**
         * Stability: 1 - Experimental
         * Added in: v13.11.0, v12.17.0
         * https://nodejs.org/api/async_context.html#asynclocalstorageenterwithstore
         */
        this.#storage.enterWith(options.defaultValue)
      }

      asyncContextSetForSnapshot.add(this)
      asyncStorageClean.register(this, this.#storage)
    }

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

    get(): T | undefined {
      const value = this.#storage.getStore()
      if (value === undefined) {
        return this.#defaultValue
      }
      return value
    }

    run<R>(value: T, fn: (...args: any[]) => R, ...args: any[]): R {
      return this.#storage.run(value, fn, ...args)
    }
  }

  export interface AsyncVariableOptions<T> {
    name?: string;
    defaultValue?: T;
  }

  export class Snapshot {
    run<R>(fn: () => R): R {
      const combined = Snapshot.wrap(fn)

      return combined()
    }

    /**
     * same like Experimental `Static method: AsyncLocalStorage.bind(fn)`
     *
     * https://nodejs.org/api/async_context.html#static-method-asynclocalstoragebindfn
     */
    static wrap<R>(fn: () => R): () => R {
      const allAsyncContexts = [...asyncContextSetForSnapshot]
      if (!allAsyncContexts.length) {
        return fn
      }

      const snapshotValues = allAsyncContexts.map(context => [context, context.get()] as const)

      const combined = snapshotValues.reduce((prevFn, [context, value]) => {
        return () => context.run(value, prevFn)
      }, fn)

      return combined
    }
  }
}

const asyncContextSetForSnapshot = new IterableWeakSet<AsyncContext.Variable<unknown>>()

const asyncStorageClean = new FinalizationRegistry<AsyncLocalStorage<unknown>>(storage => {
  /**
   * Stability: 1 - Experimental
   * Added in: v13.10.0, v12.17.0
   * https://nodejs.org/api/async_context.html#asynclocalstoragedisable
   */
  storage.disable()
})

IterableWeakSet.ts

/**
 * Basic implementation of IterableWeakSet
 *
 * like IterableWeakMap in https://github.com/tc39/proposal-weakrefs#iterable-weakmaps
 * via [WeakRef](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/WeakRef)
 */
export class IterableWeakSet<T extends WeakKey> {
  #weakMap = new WeakMap<T, {
    value: T;
    ref: WeakRef<T>;
  }>();
  #refSet = new Set<WeakRef<T>>();
  #finalizationGroup = new FinalizationRegistry<CleanGroup<T>>(IterableWeakSet.#cleanup);

  static #cleanup<T extends WeakKey>({ set, ref }: CleanGroup<T>) {
    set.delete(ref);
  }

  add(value: T) {
    const ref = new WeakRef(value);

    this.#weakMap.set(value, { value, ref });
    this.#refSet.add(ref);

    this.#finalizationGroup.register(value, {
      set: this.#refSet,
      ref
    }, ref);

    return this
  }

  delete(value: T) {
    const entry = this.#weakMap.get(value);
    if (entry === undefined) {
      return false;
    }

    this.#weakMap.delete(value);
    this.#refSet.delete(entry.ref);
    this.#finalizationGroup.unregister(entry.ref);
    return true;
  }

  *[Symbol.iterator](): Generator<T> {
    for (const ref of this.#refSet) {
      const value = ref.deref();
      if (value === undefined) continue;
      yield value;
    }
  }

  *values(): Generator<T> {
    for (const value of this) {
      yield value;
    }
  }
}

interface CleanGroup<T extends WeakKey> {
  set: Set<WeakRef<T>>;
  ref: WeakRef<T>;
}