atomFamily.remove() does not trigger re-renders for subscribed components

[taylorcode]

Bug Description

When using atomFamily.remove(param), components that are subscribed to the underlying atom do not re-render.

Currently, a workaround is to set a noop/empty value before calling remove(), but this must be done carefully, and only works when the atom has a valid empty state.

import React from 'react';
import {atom, useAtomValue, useSetAtom} from 'jotai';
import {atomFamily} from 'jotai/utils';

const itemFamily = atomFamily((id: number) =>
  atom({
    value: Math.random(),
  })
);

const itemKeysAtom = atom<number[]>([1, 2, 3]);

const resetItemsAtom = atom(null, (get, set) => {
  const keys = get(itemKeysAtom);
  for (const id of keys) {
    // Without this line, subscribed components do not re-render
    // despite the internal atom values being reset
    set(itemFamily(id), {value: 0});
    // This should trigger a re-render by itself, because it causes the
    // atom values to change, but it does not
    itemFamily.remove(id);
  }
});

function Item({id}: {id: number}) {
  const item = useAtomValue(itemFamily(id));
  return <div>{item.value}</div>;
}

export default function App() {
  const keys = useAtomValue(itemKeysAtom);
  const reset = useSetAtom(resetItemsAtom);

  return (
    <div>
      {keys.map((id) => (
        <Item key={id} id={id} />
      ))}
      <button onClick={reset}>Reset (should render new random numbers)</button>
    </div>
  );
}

Reproduction Link

https://codesandbox.io/p/sandbox/busy-aj-hgjvhx

[dai-shi]

That's the designed behavior. Creating a new atom with atomFamily doesn't trigger rerender, and neither does removing.

I think you want to look into the atom-in-atom pattern.

[taylorcode]

I find the behavior to feel bug-like. My specific issue is related to handling user sign outs - When the user signs out, all user-related state must be cleared in the UI. I had a related bug where the prior session state is still displayed in the UI, simply because nothing triggered the component to re-render when the family's param was reset. As soon as something else causes it to be re-rendered, the component will render the reset/default state, which is the real/current internal state of the atom.

I've migrated away from using arrays of atoms (I think that's what you mean by atom-in-atom) to using atomFamily everywhere for a few reasons. It provides a simpler mental model because with families, atoms are created within the state graph automatically as components are rendered, rather than being created and managed manually within application logic. It's also much like database normalization and provides similar properties/benefits.

For example with families it's easy to store metadata at the key or index level:

export const favoritesCollectionCid = 'favorites';

const favoritesCollectionKey: ICollectionKey = {
  cid: favoritesCollectionCid,
  lastModifiedTime: MAX_DATE_MS, // ensures that it is always first when sorted by last modified time
};

export const collectionKeysAtom = atomWithReset<ICollectionKey[]>([favoritesCollectionKey]);

export const collectionCidsAtom = atom((get) => {
  return get(collectionKeysAtom)
    .sort((a, b) => b.lastModifiedTime - a.lastModifiedTime)
    .map((k) => k.cid)
});

export const collectionFamily = atomFamily((_cid: string) => {
  return atom<Collection>({
    name: 'Favorites',
  });
});

(cid is short for clientId. The server also sends an id)

In this example, I'm able to sync the lastModifiedTime for all collections from my database, allowing the component that renders the list of collections to display them sorted by this property and (importantly) without subscribing to changes to the individual atoms.

[dmaskasky]

But with the hook I provided below, the code changes to this. You could treat this as a best practice.

const studentsFamily = atomFamily((param: string) => atom(param))

function Component() {
  const studentAtom = useAtomFamilyAtom(studentsFamily, 'abc')
  const value = useAtomValue(studentAtom)
  
  return <>{value}</>
}

[taylorcode]

atomFamily creates each atom and internally maps them to the param. So when calling someFamily.remove(param) it could look up the atom, and trigger its subscribing components to render using the same mechanism as set(someFamily(param), undefined), but this is just an assumption about the implementation details.

[dmaskasky]

It's possible to do this with building blocks. But I am unsure we want to have this behavior built-in.

A reference to the atom could be kept external from the atomFamily. There is no way to tell if the component that subscribes to the atom is receiving that atom from the atomFamily or from some external reference.

[taylorcode]

Right I gather that the challenge is someFamily(someParam) performs a Map lookup and returns the value, rather than observing the property internalAtomMap[someParam] for changes.

[taylorcode]

[in reply to a former question about whether remove() should unmount the atom]

I think it should. The atom is created solely by the atomFamily, and so it makes sense that it would also be deleted by it.

Otherwise the atomFamily only manages one part of its lifecycle, creation, leaving the other part to be done manually by the developer.

[dai-shi]

FYI, we plan to move atomFamily to a separate package. #2889

[dmaskasky]

You can already do this

/**
 * @returns the atom for the given param
 * @sideeffect calls setState when the returned atom is removed
 **/
function useAtomFamilyAtom<T extends string, U extends AnyAtom>(atomFamily<T, U>, param: T): U {
  const [, setState] = useState(0)
  const atom = atomFamily(param)
  useEffect(() => {
    return atomFamily.unstable_listen((e) => {
      if (e.type === 'REMOVE' && e.atom === atom) {
        setState((v) => v + 1)
      }
    })
  }, [atomFamily, atom])
  return atom;
}

[taylorcode]

Is .remove() useful without this behavior built-in? That is what I’m wondering about.

Once an atom is removed from a family atomFamily(param) no longer refers/maps to an atom and is essentially undefined, yet components are not made aware of this, and continue to display the (now stale) value.

[dmaskasky]

Yes, because unstable_listen provides the callback for this action. Jotai's api provides a clear boundary between vanilla and react interfaces. atomFamily is a vanilla util.

So if you want component reactivity bound to unstable_listen, the correct approach would be to stop doing atomFamily(param) inside a component. Doing atomFamily(param) inside a component is akin to doing store.get(atom) inside a component.

[taylorcode]

Considering this, it seems that using useAtom/useAtomValue with an atomFamily is a footgun, as it leads to this problem of stale React components when atoms are removed from the family, which is an essential practice. So, useAtomFamilyAtom is the solution here. Perhaps then, the real fix here is to update the types of useAtom/useAtomValue such that they will not accept someAtomFamily(param)?

[Wendystraite]

This caveat and hook could be included in the jotai's doc in a React-specific section.

[taylorcode]

I agree but the other Jotai hooks enforce (via parameter types) the kinds of atoms they accept / can be used with.

[dai-shi]

Is .remove() useful without this behavior built-in? That is what I’m wondering about.

remove() exists to provide a way to avoid memory leaks. Our mental model is that atomFamily is simply a Map that is somewhat aware of Jotai, but it doesn't depend on Jotai.

[taylorcode]

to me useAtom(someFamily(param)) reads as: “subscribe to an atom within the family, and when map.delete(param) / map.set(param, newAtom) occurs, let me know”

when in reality (based on my understanding, it is):
“look up the atom one-time from the family’s internal Map and subscribe to it”

[dai-shi]

In my mental model, it should be useAtomFamily(someFamily, param) or something.

when in reality (based on my understanding, it is):
“look up the atom one-time from the family’s internal Map and subscribe to it”

Yes. So, this is conceptually a suggestion for a new capability. I'm not completely against it, but this is a good time to move to a separate package: #3143 (comment)

[dmaskasky]

I see a few options here:

1. useAtomFamilyAtom(atomFamily, param) - calls component render when atomFamily.remove is called.
2. atomFamily.remove calls all atom listeners. What happens if we still have some useAtom(atomFamily(param))calls?useAtom` would be passed undefined in that case.
3. atomFamily.remove unmounts the atom. What happens if we still have some useAtom(atomFamily(param))calls?useAtom` would be passed undefined in that case.
4. Remove the ability to call atomFamily.remove. atomFamily would cause a memory leak that would not be addressable.
5. Convert atomFamily to use a WeakMap internally. remove would go away, but this would require passing object keys to atomFamily. This is a significant departure from existing api.
6. atomFamily.remove marks the atom for removal when the atom is unmounted. It doesn't remove the atom from the atomFamily while there are still subscribers.

Option 6 is the best, IMO, for maintaining existing api, preventing memory leaks, and addressing stale references in React components.

[taylorcode]

ha! that's clever. Taking a step back, the core issue is with components displaying atom values that are no longer inside of the atom family. I don't believe solution 6 addresses this, it just makes the stale value sticky for longer. The component continues to render the stale value until it is re-mounted.

[dmaskasky]

This is not the case. The atom will not even be removed from the atomFamily while it has subscribers. As long as the component continues to subscribe to this atom, the atom will not be removed. This avoids the issue entirely, and shifts control to the application logic to have that component stop listening to that atomFamily param before expecting the atom to be removed. Its an application bug if a component is listening to an atomFamily param that should no longer exist. The logic that handles the atom removal should address this.

[taylorcode]

got it, so your proposal changes behavior of .remove(param), which currently acts like Map.delete(param), to something like .bindAtomCleanupToComponentUnmount(param). Funny enough, Recoil actually had a feature for this, but it was never enabled due to an internal feature flag. Yet, in Recoil, it was set up on atom family definition, as a means of configuring its memory management policy for all atoms within the family.

If I understand correctly, it is also implied that:

1. atomFamily is responsible for creating atoms, and optionally, for removing them from memory when subscribing components have unmounted.
2. Atoms must be cleared/reset manually, and this is not the role of atomFamily.

[dmaskasky]

Yes. You got it. Feel free to review my PR if you like.

I'm not sure it will be accepted though. We are planning on moving atomFamily to 'jotai-family' package. That might happen soon actually.

[taylorcode]

the behavior makes sense to me + would have prevented a bug that I encountered

[dmaskasky]

#3145