Add atomFamily.getKeys method in order to simplify family keys management

[grzesiek-ds]

Problem statement.

atomFamily uses Map underneath. It is not possible to directly access given Map, while also, theres no method on atomFamily which would allow to get all keys from Map.

This leads to situation where in order to avoid memory-leaks on "infinite" atomFamily (family with unknown number of family members), we have to manually track keys which have been set in Map.

Suggested change

My suggestion is to simply add new method to atomFamily API:

jotai/src/vanilla/utils/atomFamily.ts

Line 22 in 6a16612

const atoms: Map<Param, [AtomType, CreatedAt]> = new Map()

...

createAtom.getKeys = () => {
  return atoms.keys();
}

Example approach with current API:

// Atom representing timeline boundaries accessible in UI. It updates eg. each 15 minutes.
const timelineBoudariesAtom = atom<TimelineBoundaries>({ min: ..., max: ... });
// Atom to track keys in bucketsAtom.
const timestampsAtom = atom<number[]>([]);
// Atom family which contains data for each bucket accessible on timeline (eg. each minute).
const bucketsAtom = atomFamily<number, number[]>((timestamp) => atom([])); 
// Atom to set values in both timestampsAtom and bucketsAtom.
const bucketSetterAtom = atom((null, ([timestamp, items]) => {
  set(bucketsAtom(timestamp), items);
  set(timestampsAtom, (currentTimestamps) => [...currentTimestamps, timestamp]));
})
// Atom to periodically clean atomFamily.
const bucketsCleanerAtom = atom((get, set) => {
  // Once min value of timeline boundaries has changed,
  // figure out which keys should be removed from atomFamily.
  const { min } = get(timelineBoundariesAtom);
  const timestamps = get(timestampsAtom);
  const timestampsToRemove = timestamps.filter((timestamp) => timestamp < min);
  // Remove keys from atomFamily.
  timestampsToRemove.forEach(bucketsAtom.remove);
  // Update list of keys defined in atomFamily. This will re-trigger bucketsCleanerAtom.
  set(timestampsAtom, timestamps.filter((timestamp) => !timestampsToRemove.includes(timestamp));
});

Example approach with new API:

const timelineBoudariesAtom = atom<TimelineBoundaries>({ min: ..., max: ... });
const bucketsAtom = atomFamily<number, number[]>((timestamp) => atom([])); 
// This atom is no-longer needed as it operates only on bucketsAtom now.
const bucketSetterAtom = atom((null, ([timestamp, items]) => {
  set(bucketsAtom(timestamp), items);
});

const bucketsCleanerAtom = atom((get) => {
  const { min } = get(timelineBoundariesAtom);
  const timestamps = bucketsAtom.getKeys();
  const timestampsToRemove = timestamps.filter((timestamp) => timestamp < min);

  timestampsToRemove.forEach(bucketsAtom.remove);
  // 'set' is no longer executed, therefore we won't trigger additional atom update.
});

Why can't I use atomFamily.setShouldRemove?

atomFamily.setShouldRemove executes only when we want to access given key, and it has no ability to read other atoms via getter.
If we have entry in atomFamily, which is no longer possible to access by UI (it is outside of timeline boundaries), theres no scenario where setShouldRemove could be called.

[dai-shi]

Well, it's technically possible. Probably call it getAllParams. Related #1912

I have a few comments about my hesitation:

• It feels like adding getAllParams misleads to bad design. Iterating all params should be avoided as much as possible.
  • On the other hand, adding removeAll would be good.
• The current implementation is already a bit complicated and I kind of regret adding .setShouldRemove.
  • Basically, my recommendation is to create your own atomFamily, instead of complicating Jotai's atomFamily.
• atomFamily is originally added to help migrating from Recoil, but the current api is far from compatible.
  • Still, I wonder if Recoil has atomFamily.getKeys.

That said, adding getAllParams is trivial in terms of its implementation as you suggest. So, if we clearly state the valid use case and warn the misusage in docs, it might be acceptable.

If you have an idea to improve setShouldRemove (even if it's breaking change), I would like to hear that.

[sowhatdoido]

It feels like adding getAllParams misleads to bad design. Iterating all params should be avoided as much as possible.

Can you expand a little further why this is the case? For context, I'm coming from this question, and I wanted to understand if exposing the map element directly was a bad idea. Map has a number of functions that you potentially would need to expose in order to get full support for the underlying object, when you could just return that structure for us to modify directly.

For instance, instead of creating a getAllParams function, use atomFamilyInstance.getMap().keys(); if you need to delete everything, atomFamilyInstance.getMap().clear().

I'm aligned with the idea that the Atom family acts like a simple wrapper around a Map, but should that mean there should be some way of accessing that underlying object, whether its returning the entire object or exposing functions one at a time to do the same task?

[dai-shi]

We could say atomFamily is a simple wrapper around a Map, and it tries to tell a certain practice.

If you were to use .getMap(), you wouldn't need atomFamily in the first place, and you could do whatever you want with a bare Map. (tbh, I'd rather consider dropping atomFamily util, and just add a usage as recipes.)

It feels to me like a use case requiring .getAllParams() isn't very "atomic". You would treat it "as a big state". I would personally try to put such a Map in an atom mapAtom = atom(new Map()) and pick item out of it atom((get) => get(mapAtom).get(key)).

That said, I'm still open for suggestions to re-implement atomFamily, if that's reasonably simple and has migration path, including adding .getAllParams(). There might be some use cases, and I would like to have several contributors to discuss on it thoroughly, like with @grzesiek-ds and @sowhatdoido .

[sowhatdoido]

I experimented with the proposed getMap() function in my project, but I ended up arriving at a similar conclusion @dai-shi. The atomFamily didn't necessarily accomplish much other than act as a check for whether or not something exists in the map, returning it if it did and creating it if it didn't. Throw on the fact that adding to the atom family isn't a reactive event, I didn't really have much use for it and ended up using atom(new Map()) instead.

---

I think what I was hoping atomFamily could do was help manage the use of a Map inside of atoms a little easier. Correct me if I'm wrong but the way you'd go about updating a map inside an atom would look something like this:

const collection = atom(
  new Map([
    ["foo", "some"],
    ["bar", "value"]
  ])
);


const [map, updateMap] = useAtom(collection);
<button onClick={() => updateMap(new Map([...map, ["baz", "new data"]]))}>Append</button> // Is this correct?

Assuming I'm right, I find that not being able to call a simple set(key, value) command made it unintuitive to interface with, as we normally just call Map.set(key, value) when dealing with stand alone maps.

I ended up writing this:

const atomMap = (baseAtom) => {
  const derivedAtom = atom(
    (get) => (key) => get(baseAtom).get(key),
    (get, set, key, value) => {
      const newMap = new Map(get(baseAtom));
      newMap.set(key, value);
      set(baseAtom, newMap);
    }
  );

  return derivedAtom;
};

const useAtomMap = (baseAtom) => {
  const atomCollectionInstance = useMemo(() => atomMap(baseAtom), [baseAtom]);
  return useAtom(atomCollectionInstance);
};

The atomMap (atomWithMap? I don't know what to call it) is a derived atom that allows me to simply access the map entries with an intuitive get and set function, and the hook just let me take some shortcuts (otherwise you have to declare the base Atom and the atomMap separately before passing it into useAtom).

  const [get, set] = useAtomMap(collection);
  get("key") // acts like Map.get('key')
  set("key", "value") // acts like Map.set("key", "value")

I found I liked this solution more because I felt like it was a little more atomic in nature. I'm able to start with a simple atom that contains a map, then pass it into an atom creator that exposes those get and set functions, while still being able to maintain a reference to the map when I need it, and all while being reactive.

  const [map] = useAtom(collection); // Can interate through this
  const [get, set] = useAtomMap(collection); // Use these to fetch and update individual values

The above isn't necessarily a proposal but more along the lines of what I was looking to do and where I ended up, just in case it helps anybody else :) Also if my code is terrible, please feel free to let me know - I'm still trying to wrap my head around jotai. Loving it so far though!

[dai-shi]

Sounds like you are going the right way. atomFamily isn't the right tool for your case.

atomWithReducer might help:

const collection = atomWithReducer(
  new Map([
    ["foo", "some"],
    ["bar", "value"]
  ]),
  (prev, action) => {
    if (action.type === 'set') {
      return new Map(prev).set(action.key, action.value)
    }
    return prev
  }
);

Re: atomMap

    (get) => (key) => get(baseAtom).get(key),

That doesn't work. You can't return a function that delays get invocation. It will not be reactive.

[sowhatdoido]

thank you for your feedback!

I've updated it to:

   (get) => {
      const map = get(baseAtom);
      return (key) => map.get(key);
    },

this will immediately invoke the get function while retaining the same functionality I was looking for.

[grzesiek-ds]

@dai-shi my suggestion for implementation changes of atomFamily, theres few breaking changes. Let me know what do You think of it.
I decided to not to expose Map keys, instead there are two different call signatures for remove method. Although, I think that it could be good idea to expose keys iterator (not array of keys).

1. Items from Map are now extracted via Map.get(key) instead of iterating over Map keys.
2. Theres 3 instead of 2 type arguments.

• Param - as previous.
• Key - it describes keys of Map, has to be plain value (string | number | boolean).
• AtomType - as previous.

3. Second argument - areEqual has been replaced with paramSerializer. paramSerializer is required when Param does not extends Key, otherwise it's optional. paramSerializer accepts Param as argument, and is expected to return Key.
4. setShouldRemove has been removed.
5. remove now has two call signatures.

• remove(param: Param) - same as before.
• remove((key: Key, abortAfter: VoidFunction) => boolean) - function which iterates over Map keys. Depending on return value (boolean) key is either removed (true) or kept (false). In order to early abort Map iteration, abortAfter function can be called.

6. Date of item creation is no longer kept in Map - but if You see use case for it - it can be added back.

import { Atom } from 'jotai';

export type RemoveFamilyParamCallback<Key> = (key: Key, abortAfter: VoidFunction) => boolean;
export type RemoveFamilyParam<Param, Key> = (removeParam: RemoveFamilyParamCallback<Key> | Param) => void;
export type ParamSerializer<Param, Key extends string | number | boolean> = (param: Param) => Key;

export interface TestAtomFamily<Param, Key, AtomType> {
  (param: Param): AtomType;
  remove: RemoveFamilyParam<Param, Key>;
}

const isRemoveFamilyParamCallback = <Key>(
  possiblyCallback: unknown
): possiblyCallback is RemoveFamilyParamCallback<Key> => typeof possiblyCallback === 'function';

export function testAtomFamily<Param, Key extends string | number | boolean, AtomType extends Atom<unknown>>(
  initializeAtom: (param: Param) => AtomType,
  ...[paramSerializer]: Param extends Key
    ? [paramSerializer?: ParamSerializer<Param, Key>]
    : [paramSerializer: ParamSerializer<Param, Key>]
): TestAtomFamily<Param, Key, AtomType> {
  const getKey = (param: Param) => (paramSerializer ? paramSerializer(param) : (param as unknown as Key));
  const atoms: Map<Key, AtomType> = new Map();

  const getAtom = (param: Param) => {
    const key = getKey(param);
    const item = atoms.get(key);

    if (item) {
      return item;
    }

    const newItem = initializeAtom(param);
    atoms.set(key, newItem);

    return newItem;
  };

  const removeFamilyParam: RemoveFamilyParam<Param, Key> = removeParam => {
    if (isRemoveFamilyParamCallback(removeParam)) {
      let shouldAbortAfter = false;
      const abortAfter = () => {
        shouldAbortAfter = true;
      };

      for (const key of atoms.keys()) {
        if (removeParam(key, abortAfter)) {
          atoms.delete(key);
        }

        if (shouldAbortAfter) {
          break;
        }
      }
    } else {
      atoms.delete(getKey(removeParam as Param));
    }
  };

  getAtom.remove = removeFamilyParam;

  return getAtom;
}

Theres stackblitz playground:
https://stackblitz.com/edit/vitejs-vite-gcxqvc?file=src%2FApp.tsx

Note! Types management should be improved.

[dai-shi]

Honestly, that's too much change. (It could be a third-party library.)
I would like to keep Param and areEqual, without key, no serialization. For performance, we could serialize in advance.

I'm okay with deprecating (then eventually remove) setShouldRemove and Date (we could encode it in Param on userland), and changing remove to accept a function (maybe it should be a new method).

[grzesiek-ds]

Will provide another attempt at the end of week.

[icambron]

My use case for this is essentially the one at in the CodeSandbox example. Essentially, I'm hoping to avoid having to keep track of a separate atom of the items, which the example does via:

const todosAtom = atom<string[]>([]);

and iterates over like:

 return todos.filter((id) => !get(todoAtomFamily({ id })).completed);

In more complex scenarios, keeping todosAtom and todosAtomFamily in sink is a pain and bug-prone.

An an alternative approach I'm considering is putting the whole map in an Atom. To iterate over it, just iterate over the value. To treat individual elements as atoms, use an optic. Probably what I really want is splitAtom but for maps instead of lists.

[dai-shi]

Can you try creating your own atomFamily with getKeys capability and share how it goes? I really wonder if it solves your use case.

splitAtom but for maps instead of lists.

For read-only atom, you can derive mapAtom -> listAtom -> splitAtom, no?

[icambron]

I'll see what I can come up with and share it back.