LRU style atomFamily

[ajoslin]

Use-case: I'm using an atomFamily to model recently fetched messages within multiple channels in a chat application. Family factory = channel id, return = array of messages.

There can be potentially hundreds of channels, so to save on memory I'd like to use an "LRU style" atomFamily where it would drop the least recently used keys after it reaches X size. Then if the user clicks every channel it'd drop the least recently viewed channel's messages from the cache.

I see that setShouldRemove does exist, but that doesn't seem to have a way to reference the size of the whole map and compare an item to the oldest item in the map.

Has anyone done this sort of thing / any pointers on how to? I suppose I could write a custom atomLruFamily

[dai-shi]

We didn't include size option in atomFamily, because it's Map-based. LRU would require array like structure. (or createdAt can do, but with the cost of iterating all atoms in the map.)

I recently made jotai-cache, which has size option.
I think writing atomLruFamily is trivial.

Refer the simplified atomFamily #1323 (reply in thread)

Here's the one based on a naive array.

const atomLruFamily = (initializeAtom, { size }) => {
  const cache = []
  return (param) => {
    const found = cache.find((item) => item[0] === param))
    if (found) return found[1]
    const atom = initializeAtom(atom)
    cache.unshift([param, atom])
    if (cache.length > size) cache.pop()
    return atom
  };
};

[ajoslin]

Awesome - thanks @dai-shi - super simple indeed:

import { PrimitiveAtom } from "jotai";

export type AtomLruFamily<Key, Value> = ((key: Key) => PrimitiveAtom<Value>) & {
  keys: () => Key[];
  onBeforeEvict?: (key: string) => void;
  map: Map<Key, PrimitiveAtom<Value>>;
};

export function atomLruFamily<Key, Value>(
  initializeAtom: (key: Key) => PrimitiveAtom<Value>,
  maxSize = 100
) {
  const cache = new Map();
  const keysSet = new Set<Key>();
  const fn: AtomLruFamily<Key, Value> = (key: Key) => {
    keysSet.delete(key);
    keysSet.add(key);
    if (keysSet.size > maxSize) {
      const firstKey = keysSet.values().next().value;
      fn.onBeforeEvict?.(firstKey);
      keysSet.delete(firstKey);
      cache.delete(firstKey);
    }

    if (!cache.has(key)) {
      cache.set(key, initializeAtom(key));
    }
    return cache.get(key);
  };

  fn.keys = () => [...keysSet.keys()];
  fn.map = cache;

  return fn;
}

[ajoslin]

Added an onBeforeEvict so I can handle unlisten/unmount logic from the outside (even when unmounted, channel state listens to new messages - until the LRU kicks it out).

and now this is working perfect 👌