request-based cache with Async Local Storage

request-cache.server.ts

import { AsyncLocalStorage, AsyncResource } from "node:async_hooks";

const UNTERMINATED = 0;
const TERMINATED = 1;
const ERRORED = 2;

type Status = typeof UNTERMINATED | typeof TERMINATED | typeof ERRORED;

type CacheNode<T> = {
  s: Status;
  v: T | undefined;
  o: WeakMap<object, CacheNode<T>> | null;
  p: Map<
    string | number | null | undefined | symbol | boolean,
    CacheNode<T>
  > | null;
};

function createCacheNode<T>(): CacheNode<T> {
  return {
    s: UNTERMINATED,
    v: undefined,
    o: null,
    p: null,
  };
}

type CacheStore = WeakMap<Function, CacheNode<any>>;

// Create an AsyncLocalStorage instance to store request-scoped data
const asyncLocalStorage = new AsyncLocalStorage<CacheStore>();

function getDebugLabel(fn: Function): string {
  return fn.name || "anonymous";
}

/**
 * Wraps a function with request-scoped caching using AsyncLocalStorage.
 * Heavily inspired by React.cache, but uses AsyncLocalStorage to store the cache
 * so that it can be used in any context (e.g. data loaders) and not just in React components.
 * Similar to React.cache, it automatically memoizes based on function arguments.
 * The wrapped function will be executed only once per request for the same arguments,
 * and subsequent calls within the same request will return the cached result.
 */
export function cache<TArgs extends unknown[], TResult>(
  fn: (...args: TArgs) => Promise<TResult>,
): (...args: TArgs) => Promise<TResult> {
  const fnName = getDebugLabel(fn);
  return async function (...args: TArgs) {
    const store = asyncLocalStorage.getStore();
    console.log(
      `[cache:${fnName}] AsyncLocalStorage store:`,
      store ? "present" : "missing",
      "at:",
      new Error().stack,
    );
    if (!store) {
      console.log(
        `[cache:${fnName}] No cache store found, executing without cache`,
      );
      return fn(...args);
    }

    let cacheNode = store.get(fn) ?? createCacheNode<TResult>();
    store.set(fn, cacheNode);

    // Traverse the cache tree based on arguments
    for (let i = 0; i < args.length; i++) {
      const arg = args[i];
      if (
        typeof arg === "function" ||
        (typeof arg === "object" && arg !== null)
      ) {
        // Objects go into a WeakMap
        if (cacheNode.o === null) {
          cacheNode.o = new WeakMap();
        }
        const nextNode = cacheNode.o.get(arg) ?? createCacheNode<TResult>();
        cacheNode.o.set(arg, nextNode);
        cacheNode = nextNode;
      } else {
        // Primitives go into a regular Map
        const primitiveArg = arg as
          | string
          | number
          | boolean
          | symbol
          | null
          | undefined;
        if (cacheNode.p === null) {
          cacheNode.p = new Map();
        }
        const nextNode =
          cacheNode.p.get(primitiveArg) ?? createCacheNode<TResult>();
        cacheNode.p.set(primitiveArg, nextNode);
        cacheNode = nextNode;
      }
    }

    if (cacheNode.s === TERMINATED) {
      console.log(`[cache:${fnName}] Cache hit for args:`, args);
      return cacheNode.v as TResult;
    }
    if (cacheNode.s === ERRORED) {
      console.log(`[cache:${fnName}] Cache hit (error) for args:`, args);
      throw cacheNode.v;
    }

    try {
      console.log(`[cache:${fnName}] Cache miss for args:`, args);
      const result = await fn(...args);
      cacheNode.s = TERMINATED;
      cacheNode.v = result;
      return result;
    } catch (error) {
      console.log(`[cache:${fnName}] Cache miss (error) for args:`, args);
      cacheNode.s = ERRORED;
      cacheNode.v = error;
      throw error;
    }
  };
}

class RequestContext extends AsyncResource {
  constructor() {
    super("RequestContext");
  }

  run<T>(store: CacheStore, fn: () => Promise<T>): Promise<T> {
    return this.runInAsyncScope(async () => {
      return asyncLocalStorage.run(store, fn);
    });
  }
}

/**
 * Creates a new request context with a fresh cache store.
 * This should be called at the beginning of each request.
 */
export async function withRequestCacheContext<T>(
  fn: () => Promise<T>,
): Promise<T> {
  console.log("[cache] Creating new request cache context");
  const requestContext = new RequestContext();
  const result = await requestContext.run(new WeakMap(), fn);
  console.log("[cache] Request cache context created");
  requestContext.emitDestroy();
  return result;
}