createStoredObjectsHook.ts

import { useCallback, useState } from "react";
import { useAsyncEffect } from "./useAsyncEffect";

/**
 * Hook factory to create hooks for storage models.
 * The resulting hook allows components to use stored objects in their state.
 *
 * Feel free to copy and paste this documentation comment into the resulting
 * hook function (replace "something" with the collection name):
 *
/**
 * (This comment was copied from the definition for createStoredObjectsHook())
 *
 * Use something in state.
 *
 * Usage:
 *
 *   const [
 *     something,
 *     updateSomething,
 *     createSomething,
 *     deleteSomething
 *   ] = useSomething(...);
 *
 * You can pass in:
 *
 * - nothing to get all objects
 * - a list of object IDs
 * - a custom fetcher function that returns objects to hold in state
 *
 * Subscribes to changes in returned objects, such that if any other component
 * updates one of them it triggers a rerender.
 *
 * Returns null initially to signal a loading state.
 *
 * Objects are cached in the following way:
 *
 * - Hooks passing nothing (thus fetching all objects) and those passing a
 *   fetcher function always read from the database on the first render to
 *   obtain a list of object IDs from the result. On subsequent renders they
 *   try to read from cache first.
 * - Hooks passing a list of object IDs always try to read from cache before
 *   fetching from the database.
 * - The cache can have a max age for objects, after which the object is
 *   refetched on its next retrieval attempt. Check implementation for details.
 */
export function createStoredObjectsHook<
  T extends Record<string, unknown>,
  C extends unknown[],
  U extends unknown[],
>(
  storageFunctions: {
    getId: (obj: T) => string;
    create: (...args: C) => Promise<T>;
    getAll: () => Promise<Record<string, T>>;
    get: (...ids: string[]) => Promise<Record<string, T>>;
    update: (...args: U) => Promise<T>;
    delete: (...ids: string[]) => Promise<void>;
  },
  { cacheMaxAgeSeconds = 5 * 60 }: { cacheMaxAgeSeconds?: number } = {},
) {
  const cache = new Cache<T>(cacheMaxAgeSeconds);
  const pubsub = new PubSub<T>((...ids) => cache.tryGet(...ids));

  // Clean cache every minute.
  setInterval(() => {
    cache.collectGarbage();
  }, 60 * 1000);

  type Return = [
    // The objects
    Record<string, T> | null,
    // Update
    (...args: U) => Promise<void>,
    // Create
    (...args: C) => Promise<void>,
    // Delete
    typeof storageFunctions.delete,
  ];

  // Overload for fetching all objects.
  function useStoredObjects(): Return;

  // Overload for fetching from a list of object IDs.
  function useStoredObjects(objectIds: string[]): Return;

  // Overload for providing a custom fetch function that returns the objects to
  // hold in state.
  function useStoredObjects(fetchObjects: () => Promise<Record<string, T>>): Return;

  function useStoredObjects(
    fetcherOrObjectIds?: (() => Promise<Record<string, T>>) | string[],
  ): Return {
    const [objects, setObjects] = useState<Record<string, T> | null>(null);

    const fetchObjects = typeof fetcherOrObjectIds === "function" ? fetcherOrObjectIds : undefined;
    const objectIds = Array.isArray(fetcherOrObjectIds) ? fetcherOrObjectIds : undefined;

    // Initial data loading.
    // This is only meant to run once, as it subscribes to changes in the
    // returned objects.
    useAsyncEffect(
      async (signal) => {
        let objs: Record<string, T>;
        let freshObjs: Record<string, T>;

        if (objectIds) {
          const cached = cache.tryGet(...objectIds);

          const nonCachedIds = objectIds.filter((id) => !(id in cached));
          freshObjs = await storageFunctions.get(...nonCachedIds);

          objs = {
            ...cached,
            ...freshObjs,
          };
        } else if (fetchObjects) {
          objs = await fetchObjects();
          freshObjs = objs;
        } else {
          // Fetch all objects.
          objs = await storageFunctions.getAll();
          freshObjs = objs;
        }

        cache.update(freshObjs);

        if (signal.aborted) return;

        // This pub may not trigger anything if all objects were already
        // cached, so we still call setObjects() below.
        pubsub.pub(Object.keys(freshObjs));

        pubsub.sub(setObjects, Object.keys(objs));
        setObjects(objs);
      },
      [fetchObjects, objectIds?.join(",") ?? ""],
      () => {
        const orphanIds = pubsub.unsub(setObjects);
        cache.delete(...orphanIds);
      },
    );

    /**
     * Each wrapper should:
     *
     * 1. do its action
     * 2. subscribe to new IDs (if relevant)
     * 3. update cache (if relevant)
     * 4. publish update to subscribers (if relevant)
     */

    const updateWrapper = useCallback(async (...args: U) => {
      const updated = await storageFunctions.update(...args);
      const id = storageFunctions.getId(updated);

      cache.update({ [id]: updated });

      pubsub.pub([id]);
    }, []);

    const createWrapper = useCallback(async (...args: C) => {
      const newObj = await storageFunctions.create(...args);
      const id = storageFunctions.getId(newObj);

      pubsub.sub(setObjects, [id]);

      cache.update({ [id]: newObj });

      pubsub.pub([id]);
    }, []);

    const deleteWrapper = useCallback(async (...ids: string[]) => {
      await storageFunctions.delete(...ids);

      // Cache is not guaranteed to contain all objects components are listening
      // for.
      cache.delete(...ids);

      // Listeners will not get deleted objects as they are no longer in cache.
      pubsub.pub(ids);

      // Finally remove the ID from the pubsub registry.
      // We couldn't do this before triggering those listeners.
      pubsub.unsubAll(ids);
    }, []);

    return [objects, updateWrapper, createWrapper, deleteWrapper];
  }

  return useStoredObjects;
}

/**
 * Simple ID based cache with an optional max age.
 */
class Cache<T> {
  private cache: Record<string, [T, number]> = {};

  constructor(private maxAgeSeconds = -1) {}

  private isExpired(timestamp: number) {
    return this.maxAgeSeconds >= 0 && Date.now() - timestamp > this.maxAgeSeconds * 1000;
  }

  update(objs: Record<string, T>) {
    Object.entries(objs).forEach(([id, obj]) => {
      this.cache[id] = [obj, Date.now()];
    });
  }

  /**
   * Tries to get objects given a set of IDs.
   * Only those that are present in cache are returned.
   */
  tryGet(...ids: string[]) {
    const result: Record<string, T> = {};

    for (const id of ids) {
      if (!(id in this.cache)) {
        continue;
      }

      const [item, timestamp] = this.cache[id]!;

      if (this.isExpired(timestamp)) {
        delete this.cache[id];
        continue;
      }

      result[id] = item;
    }

    return result;
  }

  delete(...ids: string[]) {
    ids.forEach((id) => {
      delete this.cache[id];
    });
  }

  collectGarbage() {
    Object.entries(this.cache).forEach(([key, [_, timestamp]]) => {
      if (this.isExpired(timestamp)) {
        delete this.cache[key];
      }
    });
  }
}

type Sub<T> = (objs: Record<string, T>) => void;

/**
 * Simple pubsub class that allows:
 *
 *   - functions to register as subscribers to sets of IDs
 *   - triggering those subscribers by providing a record mapping ids to values
 */
class PubSub<T> {
  // We keep two data structures with the same data for more efficient lookups.
  private idsPerSub = new Map<Sub<T>, string[]>();
  private subsPerId: Record<string, Set<Sub<T>>> = {};

  constructor(private getter: (...ids: string[]) => Record<string, T>) {}

  sub(sub: Sub<T>, ids: string[]) {
    if (ids.length === 0) return;

    ids.forEach((id) => {
      if (!(id in this.subsPerId)) {
        this.subsPerId[id] = new Set();
      }

      this.subsPerId[id]!.add(sub);
    });

    this.idsPerSub.set(sub, [...(this.idsPerSub.get(sub) ?? []), ...ids]);
  }

  /**
   * Returns a list of IDs of objects this sub was the last subscriber of,
   * which is useful for cache eviction.
   */
  unsub(sub: Sub<T>) {
    const ids = this.idsPerSub.get(sub);

    if (!ids) return [];

    const orphanIds: string[] = [];

    this.idsPerSub.delete(sub);

    ids.forEach((id) => {
      const subs = this.subsPerId[id];

      if (!subs) return;

      subs.delete(sub);

      if (subs.size === 0) {
        delete this.subsPerId[id];

        orphanIds.push(id);
      }
    });

    return orphanIds;
  }

  pub(ids: string[]) {
    const subs = new Set<Sub<T>>();

    ids.forEach((id) => {
      this.subsPerId[id]?.forEach((x) => {
        subs.add(x);
      });
    });

    subs.forEach((sub) => {
      const subIds = this.idsPerSub.get(sub);

      if (subIds) sub(this.getter(...subIds));
    });
  }

  /** Unsubscribe all subscribers to the given IDs and delete the IDs. */
  unsubAll(ids: string[]) {
    ids.forEach((id) => {
      if (id in this.subsPerId) {
        delete this.subsPerId[id];
      }
    });

    Array.from(this.idsPerSub.keys()).forEach((sub) => {
      const newIds = this.idsPerSub.get(sub)!.filter((x) => !ids.includes(x));

      if (newIds.length === 0) {
        this.idsPerSub.delete(sub);
      } else {
        this.idsPerSub.set(sub, newIds);
      }
    });
  }
}