Understanding how jotai's state management works

Recently I have been looking at source codes of some good open source libraries and every time I be like, Wow.. how didn't I think about something like this yet.

Today we are going to talk about jotai, an atomic state management solution for React JS. The creator of jotai Daishi Kato, recently shared a tweet about simplified version of jotai. This article will be around that tweet. It was quite interesting how simple thing can make huge impact on React Ecosystem.

Let's understand this.

Copycopy code to clipboard
1// This is how we create atoms in react codebase using jotai
2import  { atom }  from  'jotai'
3
4const messageAtom =  atom('hello')
5const countAtom = atom(1)

State in jotai is just a collection of atoms. Atoms are not tied with any specific component, that means you can export this messageAtom and import it any other module. This is how it is implemented.

Copycopy code to clipboard
1export const atom = (initialValue) => ({ init: initialValue });

It's a function that returns a config object. This object holds the initial value of that atom.

Okay, we might have many atoms like this in our codebase, right ? So now we need some map kind of thing to hold them all. Here comes the atomStateMap.

Copycopy code to clipboard
1const atomStateMap = new WeakMap();

Why WeakMap though ? It helps in garbage collection. Key in WeakMap can only be an object and it doesn't prevent garbage collection of key objects.

Copycopy code to clipboard
1let john = { name: "John" };
2let weakMap = new WeakMap();
3weakMap.set(john, "...");
4john = null; // overwrite the reference
5
6// john is removed from memory!
7weakmap.get(john) // undefined

If reference of john doesn't exist, the key is removed from weakmap and free it's value if not referenced anywhere else.

Get and Set atoms in atomStateMap.

Copycopy code to clipboard
1const getAtomState = (atom) => {
2    let atomState = atomStateMap.get(atom);
3    if(!atomState) {
4        atomState = { value: atom.init, listeners: new Set() };
5        atomStateMap.set(atom, atomState);
6    }
7    return atomState;
8}

getAtomState returns the atom's state from atomStateMap. If there is no entry of that atom, it registers that atom in atomStateMap. atomState contains current value of that atom and a set of listeners. Listeners are simply a set of callback functions for specific components.

Let's say in this example, IncrementCounter and DecrementCounter are using the same atom. So if value of countAtom is changed in a component then we need to notify other component about that state change and trigger re-rendering for both component.

Now how we are going to use this atom in our functional component ? For that we will need a custom hook called useAtom.

Copycopy code to clipboard
1export const useAtom = (atom) => {
2    const atomState = getAtomState(atom);
3    const [value, setValue] = useState(atomState.value);
4    
5    useEffect(() => {
6        const callback = () => setValue(atomState.value);
7        atomState.listeners.add(callback);
8        callback();
9        return () => atomState.listeners.delete(callback);
10    }, [atomState]);
11    
12    const setAtom = (nextValue) => {
13        atomState.value = nextValue;
14        atomState.listeners.forEach((l) => l());
15    };
16    return [value, setAtom];
17};

First we fetch the latest state of that atom via calling getAtomState function. We need to trigger re-rendering if atom's value changes, right ? so what's better than using a useState hook.

Now we need to register our subscription for atom's value change. We create a callback function and push it in listeners array. We should unsubscribe that subscription if component is unmounted.

setAtom provides us functionality to update atom's value. Here we calling each callback in our listeners array. So it will trigger re-rendering for each of that component.

What happens if components that are using this atom are in different routes ? Since atom is not attached with any component, getAtomState will fetch the latest atom's state (that might have updated in some other pages) and register a subscription for that atom.

A simple example of useAtom hooks

Copycopy code to clipboard
1const countAtom = atom(0);
2
3const Counter = () => {
4  const [count, setCount] = useAtom(countAtom);
5  const inc = () => setCount(count + 1);
6  return (
7    <div>
8      {count} <button onClick={inc}>+1</button>
9    </div>
10  );
11};

Ending Notes

That's all. Now you understand the magic behind jotai's state management, on a simplified level of course. This might help you in contributing jotai's open-source repository or just for understanding how global statement management work under the hood.