@beignet/core/src/memo/index.ts

Core framework primitives for Beignet

/**
 * @beignet/core/memo
 *
 * Request-scoped memoization for port lookups and other async functions.
 *
 * `createMemo(fn)` deduplicates calls with the same arguments for the
 * lifetime of one request: the server enters a memo scope around every HTTP
 * request and `runServiceContext(...)` execution, and the scope's cache dies
 * with it. There is no TTL and no cross-request state, so memoized reads can
 * never serve data staler than the request that fetched them.
 *
 * Outside any scope — plain scripts, `createServiceContext(...)` callers —
 * memoized functions call straight through without caching.
 */

import { AsyncLocalStorage } from "node:async_hooks";

/**
 * Error thrown when a default cache key cannot be derived from arguments.
 */
export class MemoKeyError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "MemoKeyError";
  }
}

/**
 * Instrumentation event emitted by memoized functions inside a scope that
 * carries a recorder.
 */
export interface MemoInstrumentationEvent {
  /** Whether the call was served from the scope cache or filled it. */
  kind: "hit" | "miss";
  /** Memo name from `createMemo` options, defaulting to the function name. */
  memo: string;
  /** Encoded argument key for the call. */
  key: string;
  /** Fill duration for misses, rounded to two decimals. */
  durationMs?: number;
  /** Set when a miss's underlying call rejected (the entry is evicted). */
  failed?: boolean;
}

type MemoScopeStore = {
  entries: Map<string, unknown>;
  record?: (event: MemoInstrumentationEvent) => void;
};

const memoScope = new AsyncLocalStorage<MemoScopeStore>();

/**
 * Run a function inside a memo scope with an explicit recorder.
 *
 * Internal to the server runtime — apps should use `runMemoScope(...)`.
 * Always uses `AsyncLocalStorage.run` (never `enterWith`), so callers'
 * continuations never resume through a dangling frame.
 */
export function runWithMemoScope<T>(
  options: { record?: (event: MemoInstrumentationEvent) => void },
  fn: () => T,
): T {
  return memoScope.run({ entries: new Map(), record: options.record }, fn);
}

/**
 * Run a function inside a fresh memo scope.
 *
 * Use this in scripts and unit tests to give memoized functions a cache
 * lifetime. Nested scopes start empty; an active recorder is inherited so
 * instrumentation keeps flowing.
 */
export function runMemoScope<T>(fn: () => T): T {
  return runWithMemoScope({ record: memoScope.getStore()?.record }, fn);
}

/**
 * Memoized function returned by `createMemo(...)`.
 */
export type Memoized<F extends (...args: never[]) => unknown> = F & {
  /**
   * Drop the current scope's entry for these arguments, so the next call
   * re-executes. Call this from mutation methods that make a memoized read
   * stale within the same request.
   *
   * @returns Whether an entry existed.
   */
  invalidate(...args: Parameters<F>): boolean;
  /**
   * Drop all of this function's entries in the current scope.
   */
  clear(): void;
};

/**
 * Options for `createMemo(...)`.
 */
export interface CreateMemoOptions<F extends (...args: never[]) => unknown> {
  /**
   * Name used in instrumentation events and key errors. Defaults to the
   * wrapped function's name.
   */
  name?: string;
  /**
   * Build the cache key from the call arguments. Defaults to a structural,
   * type-tagged encoding of the arguments that throws `MemoKeyError` for
   * values it cannot encode deterministically (functions, symbols, class
   * instances, circular references).
   */
  key?: (...args: Parameters<F>) => string;
}

type EncodedMemoValue = readonly [string, ...unknown[]];

function describeArgType(value: unknown): string {
  if (value === null) return "null";
  if (typeof value !== "object") return typeof value;
  const ctor = (value as object).constructor?.name;
  return ctor && ctor !== "Object" ? `class instance (${ctor})` : "object";
}

function encodeMemoValue(
  value: unknown,
  memoName: string,
  position: number,
  seen: WeakSet<object>,
): EncodedMemoValue {
  if (value === null) return ["null"];

  switch (typeof value) {
    case "string":
      return ["string", value];
    case "number":
      return ["number", Object.is(value, -0) ? "-0" : String(value)];
    case "boolean":
      return ["boolean", value ? "true" : "false"];
    case "undefined":
      return ["undefined"];
    case "bigint":
      return ["bigint", value.toString()];
    case "object":
      break;
    default:
      throw new MemoKeyError(
        `[Beignet memo] Cannot build a default cache key for "${memoName}": argument ${position} is a ${typeof value}. Pass options.key to compute keys for these arguments.`,
      );
  }

  if (value instanceof Date) {
    return ["date", value.toISOString()];
  }

  if (seen.has(value as object)) {
    throw new MemoKeyError(
      `[Beignet memo] Cannot build a default cache key for "${memoName}": argument ${position} contains a circular reference. Pass options.key to compute keys for these arguments.`,
    );
  }

  if (Array.isArray(value)) {
    seen.add(value);
    const encoded: EncodedMemoValue = [
      "array",
      value.map((item) => encodeMemoValue(item, memoName, position, seen)),
    ];
    seen.delete(value);
    return encoded;
  }

  const proto = Object.getPrototypeOf(value);
  if (proto !== null && proto !== Object.prototype) {
    throw new MemoKeyError(
      `[Beignet memo] Cannot build a default cache key for "${memoName}": argument ${position} is a ${describeArgType(value)}. Pass options.key to compute keys for these arguments.`,
    );
  }

  seen.add(value as object);
  const record = value as Record<string, unknown>;
  const encoded: EncodedMemoValue = [
    "object",
    Object.keys(record)
      .sort()
      .map((key) => [
        key,
        encodeMemoValue(record[key], memoName, position, seen),
      ]),
  ];
  seen.delete(value as object);
  return encoded;
}

function defaultArgsKey(memoName: string, args: readonly unknown[]): string {
  return JSON.stringify(
    args.map((arg, index) =>
      encodeMemoValue(arg, memoName, index, new WeakSet()),
    ),
  );
}

function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
  return (
    value !== null &&
    (typeof value === "object" || typeof value === "function") &&
    typeof (value as { then?: unknown }).then === "function"
  );
}

function roundDuration(startedAt: number): number {
  return Math.round((performance.now() - startedAt) * 100) / 100;
}

let nextMemoInstance = 0;

/**
 * Wrap an async lookup so calls with the same arguments run once per request.
 *
 * Within a memo scope the first call executes and every later call with the
 * same key returns the same value — including the same in-flight promise, so
 * concurrent calls share one execution. Rejected promises are evicted, so a
 * failed lookup is retried by the next call rather than memoized. Outside a
 * scope the function calls through uncached.
 *
 * Memoize reads, not mutations. When a mutation in the same infra module
 * makes a memoized read stale, call `memoized.invalidate(...)` with the
 * read's arguments.
 *
 * @param fn - Function to memoize. `this` is not forwarded.
 * @param options - Optional name and key builder.
 * @returns The function with `invalidate(...)` and `clear()` attached.
 */
export function createMemo<F extends (...args: never[]) => unknown>(
  fn: F,
  options: CreateMemoOptions<F> = {},
): Memoized<F> {
  const name = options.name ?? (fn.name || "anonymous");
  // The map key is namespaced by instance, not name, so two memos that share
  // a name can never read each other's entries. The unit separator keeps
  // prefixes unambiguous: instance 1 never prefix-matches instance 11.
  nextMemoInstance += 1;
  const instancePrefix = `${nextMemoInstance}\u001f`;
  const keyOf = (args: Parameters<F>): string =>
    options.key ? options.key(...args) : defaultArgsKey(name, args);

  const memoized = ((...args: Parameters<F>): unknown => {
    const store = memoScope.getStore();
    if (!store) return fn(...args);

    const argsKey = keyOf(args);
    const mapKey = instancePrefix + argsKey;
    if (store.entries.has(mapKey)) {
      store.record?.({ kind: "hit", memo: name, key: argsKey });
      return store.entries.get(mapKey);
    }

    const startedAt = performance.now();
    const value = fn(...args);
    store.entries.set(mapKey, value);

    if (isPromiseLike(value)) {
      Promise.resolve(value).then(
        () => {
          store.record?.({
            kind: "miss",
            memo: name,
            key: argsKey,
            durationMs: roundDuration(startedAt),
          });
        },
        () => {
          if (store.entries.get(mapKey) === value) {
            store.entries.delete(mapKey);
          }
          store.record?.({
            kind: "miss",
            memo: name,
            key: argsKey,
            durationMs: roundDuration(startedAt),
            failed: true,
          });
        },
      );
    } else {
      store.record?.({
        kind: "miss",
        memo: name,
        key: argsKey,
        durationMs: roundDuration(startedAt),
      });
    }

    return value;
  }) as Memoized<F>;

  memoized.invalidate = (...args: Parameters<F>): boolean => {
    const store = memoScope.getStore();
    if (!store) return false;
    return store.entries.delete(instancePrefix + keyOf(args));
  };

  memoized.clear = (): void => {
    const store = memoScope.getStore();
    if (!store) return;
    for (const key of [...store.entries.keys()]) {
      if (key.startsWith(instancePrefix)) {
        store.entries.delete(key);
      }
    }
  };

  return memoized;
}