A State Management Tour: Atomic State With Jotai

A while ago, I started a journey into state management. All of this started as I was struggling with some prop drilling and Redux boilerplate. I wanted to find a simpler state management approach.

I've already explored a proxy-based approach with Valtio (see A State Management Tour: Proxy State With Valtio) and I wanted to try another approach: atomic state management.

And there I go with Jotai.

The Problem With useState and useContext

React comes with several ways to manage the state. The most common one is to use useState, but we cannot share state between components, leading to props drilling. To solve this, we can use useContext. But the main pain point is that we cannot make granular rendering with this approach. If the context changes, all the components using this context will re-render, even if just a little part of the context has changed.

To illustrate, here is an example with two components, one displaying a count A and a button to increment it, and another one displaying another count B and two buttons to increment each count. I've added a visual hint on components that are re-rendered.

Lets implement it with React.Context, try to click on the different buttons and see what is re-rendered:

We see that the entire component tree is re-rendered whatever the button we click on. This is because all the components are listening to the same context.

We can do better with Atomic State Management.

What Is Atomic State?

What if I could have a state that is independent of components, that can be defined close to the component responsible for it, and that can be shared with other components, without using React context?

It would look like this (it's pseudo code, don't try to run it, it won't work...):

// ComponentA.jsx

// Declaring my "atomic state" close to the component that is responsible for it
// And export it to be used by other components
export const countState = useState(0);

export const MyComponentA = () => {
    const [count, setCount] = countState;
    return (
        <div>
            <button onClick={() => setCount(count + 1)}>Increment</button>
            <p>{count}</p>
        </div>
    )
}

// ComponentB.jsx

// Import the atomic state
import countState from './ComponentA';

export const MyComponentB = () => {
    const [count, setCount] = countState;
    return (
        <div>
            <button onClick={() => setCount(count + 1)}>Increment A</button>
            <p>{count}</p>
        </div>
    )
}

This way, I could use the count state in another component, while keeping it as close as possible to the component that is responsible for it. And if this state changes, only the components using it will re-render.

The count state is atomic in the sense that it is not a part of a bigger state. It is a state on its own, and it can be shared between components.

Jotai In Action

Jotai allows us to create atomic states that can be shared between components.

Jotai's main concept is the atom. An atom is an immutable object state that can be shared between components. It is created using the atom function, which takes a value and an optional function that can interact with the atom's value.

So I can turn my pseudo-code into real code like this:

-const [count, setCount] = useState(0);
+// create an atom with a value of 0 and a function to set the atom's value
+export const countAAtom = atom(
+    0,
+    (get, set, value) => { set(countAAtom, value); }
+);

export const MyComponentA = () => {
-    const [count, setCount] = countState;
+    // useAtom returns the value and the function that we defined earlier
+    const [count, setCount] = useAtom(countAAtom);
    return (
        <div>
            <button onClick={() => setCount(count + 1)}>Increment</button>
            <p>{count}</p>
        </div>
    )
}

To access an atom, we use the useAtom hook. It has a very similar syntax to the useState hook, returning an array containing the value and a setter.

Now we can use this countAAtom in another component:

[...]
const countBAtom = atom(0, (get, set) => {
    set(countBAtom, get(countBAtom) + 1);
});
export const MyComponentB = () => {
    const [, setCountA] = useAtom(countAAtom);
    const [countB, setCountB] = useAtom(countBAtom);
    return (
        <div>
            <p>B : {countB}</p>
            <button onClick={() => setCountA(countA + 1)}>Increment A</button>
            <button onClick={() => setCountB(countB + 1)}>Increment B</button>
        </div>
    )
}

The Shared Atoms Pitfall

But... wait... if I run the above code, it seems to rerender everything when I click on the "IncrementA" button.

This is one of the tricky things you have to remember with Jotai:

What is atomic, is the atom, not the state nor the setter.

In this example we use the same countAAtom in component A, to use the countA, and in component B, to use the setCountA. What happens is:

1. From component B we update the value of the countAAtom,
2. The state of the atom is updated, so the atom is updated
3. And so we re-render the components linked to this Atom (components A AND B).

To avoid this, we have to use a different atom for the state and the setter. Jotai lets us do this by creating a new atom that will update the first one.

-export const countAAtom = atom(
-    0,
-    (get, set, value) => { set(countAAtom, value); }
-);
+ export const countAAtom = atom(0);
+ 
+ export const setCountAAtom = atom(null, (get, set, value) => {
+    // we access the countAAtom with the get function, and update it with the set function.
+    set(countAAtom, get(countAAtom) + 1);
+ });

This way, the setCountAAtom is not updated when the countAAtom is updated, and so the component is not re-rendered.

1. From component B we update the value of the countAAtom (via the setCountAAtom atom),
2. The state of the atom is updated, so the atom countAAtom is updated,
3. And so we re-render only the components linked to this Atom (component A). ComponentB is no longer linked to the countAAtom so it is not re-rendered.

Let's try by yourself:

Once you understand this, you understand Jotai.

Jotai With Extras

Jotai comes with some utilities to ease your work. You can, for example, use atomWithStorage to persist your atom's value in the local storage, or use splitAtom to split an array in atoms.

It also integrates well with some existing tools like Immer with jotai-immer, allowing it to handle complex nested states.

If you're into SSR (Server Side Rendering), Jotai comes with the useHydrateAtoms hook that allow you to hydrate your atoms on the server side, you can find examples in the Next.js section of jotai.

Conclusion

Jotai is a simple yet powerful tool to manage your state. Atomic state has to be well understood, but once you've gotten into some trap and figured it out, you should be OK with it.

When should you use Jotai? It's more a question of how comfortable you are with the mental model of atoms. Atoms encourage you to think about your state in a more granular way and to keep it close to the components that are responsible for it. Compared to Valtio for example (see the previous article), Jotai has less "magic" behind it, it does not rely on Proxys, and is more "React-like" in its approach. Compared to "flux" architecture libraries like Redux, Jotai is less verbose and its mental model is simpler.

In short, if you are looking for a state management that is granular, easy to use and that is not a flux architecture, Jotai is a good choice.

And if you are concerned about the possible lack of maintainability of Jotai, its popularity is in constant growth, and it is well maintained, so, go try it!