What the heck is atomic state management?

In 2020, I stepped away from building on the web to build PC game UIs using React. But I missed the web so I came back 2 years later and boy, had things changed. There were dozens of new libraries, frameworks and tools in the React space.

As I went through the recommendations from the community, I kept coming across something called "atomic state management". Apparently, it was a new way to manage state that had caught on during my time away. It sounded interesting but I wondered how it was really like in practice. In this month's post, I try to answer that question by refactoring Timo to use Jotai.

A quick introduction to Jotai and atomic state management

The idea of atomic state management was first introduced by a library from Facebook called Recoil. Jotai took this idea and ran with it by further simplifying the API. Similar to the story of Flux and Redux, Recoil struggled to rise above experimental status while Jotai kept growing. Eventually, Recoil stopped being actively maintained and the community started switching to Jotai. It is now the most mature atomic state management library.

At its core, atomic state management involves storing state in small objects called atoms. These atoms can then be imported by components to read and update the state stored within it. Updating the state of an atom also causes all components that use it to update. Here is what this looks like with Jotai.

import { atom, useAtom } from 'jotai';

// Store the value 0 in an atom
// Atoms can store any data type including objects and arrays
const countAtom = atom(0);

const Counter = () => {
  // Use the useAtom hook to read and update the countAtom
  const [count, setCount] = useAtom(countAtom);

  return (
    <div>
      Count: {count}
      <button onClick={() => setCount((count) => count + 1)}>Add one</button>
    </div>
  );
}

const SecondCounter = () => {
  // When the countAtom is changed in the Counter component, this component also re-renders 
  // with the updated value. And vice versa.
  const [count, setCount] = useAtom(countAtom);

  return (
    <div>
      Count: {count}
      <button onClick={() => setCount((count) => count + 1)}>Add one</button>
    </div>
  );
}

Using an atom with the useAtom hook is similar to using React's useState hook. So you can think of Jotai and atomic state management as useState on steroids.

However, this analogy starts to fall apart with Jotai's derived atoms. These are atoms that use values from other atoms. Here is what that looks like.

import { atom, useAtom } from 'jotai';

// The same countAtom as before
const countAtom = atom(0);

// A derived atom
const doubleCountAtom = atom(
  // The read function that is invoked when this atom is read
  // It can use get to read any other atom
  // When atoms that are read using get update, this atom recalculates its value and triggers an update
  (get) => (get(countAtom) * 2),
  // The write function that is invoked when this atom is written to
  // It can use get to read and set to write another atom.
  (get, set, value) => {
    set(countAtom, value / 2)
  }
)

// The same counter component as before
const Counter = () => {
  const [count, setCount] = useAtom(countAtom);

  return (
    <div>
      Count: {count}
      {/* When count is changed here, the value returned by the doubleCountAtom also changes */}
      <button onClick={() => setCount((count) => count + 1)}>Add one</button>
    </div>
  );
}

const DoubleCounter = () => {
  // Use the doubleCountAtom
  const [doubleCount, setDoubleCount] = useAtom(doubleCountAtom);

  return (
    <div>
      {/* doubleCount is always 2x of count due to doubleCountAtom's read function */}
      Double count: {doubleCount}
      {/* Setting doubleCountAtom's value also changes the countAtom's value due its set function */}
      <button onClick={() => setDoubleCount((doubleCount) => doubleCount + 2)}>Add two</button>
    </div>
  );
}

Derived atoms feel like magic, especially when many are chained together. Derived atoms can be made read-only or write-only by removing the write or read function. Jotai also provides hooks for accessing only the value and the update function of the atom.

// useAtomValue returns only the value of an atom
// useSetAtom returns only the function to set the value of an atom
import { useAtomValue, useSetAtom } from 'jotai';

In addition to atoms and hooks, Jotai also has utilities and extensions. The utilities add extra functionality to atoms (e.g. atomWithStorage persists the value of an atom to local storage). The extensions allow you to interact with other libraries using atoms. (e.g. Query allows you to use Tanstack Query through atoms).

Jotai in practice

To get a sense of Jotai and atomic state management in practice, I refactored Timo to use Jotai. I started with the Tanstack query version of Timo and the refactoring took me much longer than I thought. (You can find the completed refactor here)

First impressions

My first impression of Jotai was not good. To kick things off, I started replacing the user context with a Jotai atom. The user context contained the user's details and several functions for manipulating them.

// Poor man JS's enum for the user's authentication status
const UserStatus = {
  UNKNOWN: 'unknown',
  AUTHENTICATED: 'authenticated',
  UNAUTHENTICATED: 'unauthenticated'
};

// The shape of the value in the user context
const user = {
  status: UserStatus.UNKNOWN || UserStatus.AUTHENTICATED || UserStatus.UNAUTHENTICATED,
  data: {
    id: 1,
    username: 'bash',
    avatar_character: 'b',
    avatar_background: 'dark'
  },
  setAuthenticatedUser: () => {},
  clearUser: () => {}
}

The UserContextProvider component also contained the logic for fetching the user data from the API and setting the user's authentication status. I should have refactored this in the Tanstack query version to use Tanstack query but I accidentally missed it. So it still uses a simple useEffect to fetch the data.

const [user, setUser] = useState({
  status: UserStatus.UNKNOWN
});

useEffect(() => {
  if (user.status === UserStatus.UNKNOWN) {
    getUser().then((remoteUser) => {
      setUser(user => ({ ...user, status: UserStatus.AUTHENTICATED, data: remoteUser }));
    }).catch(() => {
      setUser(user => ({ ...user, status: UserStatus.UNAUTHENTICATED }));
    });
  }
}, [user]);

Since Jotai has an extension for Tanstack query (jotai-tanstack-query), I thought I could rework the user context using that extension to create an atom. This is where the struggle began.

Immediately, I got a cryptic error message when I tried to use atomWithQuery from the query extension. As it turns out, while regular atoms can accept plain objects, atomWithQuery atoms cannot.

import { atomWithQuery } from 'jotai-tanstack-query';
import { getUser } from '@timo/common/api';

/* ❌
Uncaught TypeError: t is not a function
  at Object.read (jotai-tanstack-query.js?v=07952fb4:72:32)
  at readAtomState (chunk-E7E2OHPZ.js?v=4b892653:278:36)
  at getter (chunk-E7E2OHPZ.js?v=4b892653:235:22)
  at jotai-tanstack-query.js?v=07952fb4:73:20
  at Object.read (jotai-tanstack-query.js?v=07952fb4:41:27)
  at readAtomState (chunk-E7E2OHPZ.js?v=4b892653:278:36)
  at getter (chunk-E7E2OHPZ.js?v=4b892653:235:22)
  at Object.read (jotai-tanstack-query.js?v=07952fb4:46:27)
  at readAtomState (chunk-E7E2OHPZ.js?v=4b892653:278:36)
  at getter (chunk-E7E2OHPZ.js?v=4b892653:235:22)
*/
const getUserAtom = atomWithQuery({
  queryKey: ['user'],
  queryFn: getUser,
});

// ✅
const getUserAtom = atomWithQuery(() => ({
  queryKey: ['user'],
  queryFn: getUser,
}));

Next, I tried to create a derived atom on top of the getUserAtom to create an atom with a similar shape to the old user context value. Here I ran into another head-scratcher. Previously, components could change the user context value by calling the setAuthenticatedUser function that was in it. But atomWithQuery atoms are read-only so new values can't be written to them. So how could the data in atomWithQuery be updated by other components?

I realized that the only way I can update the user data is to force getUserAtom to refetch the data. After some tinkering, I created the following userAtom where components can trigger a refetch by writing an action to the atom.

const UserAtomActions = {
  Refresh: 'refresh',
};

const userAtom = atom(
  (get) => {
    const { error, data, isFetching } = get(getUserAtom);
    if (error?.message === 'Authentication required') {
      return {
          status: UserStatus.UNAUTHENTICATED
      };
  }
    // Stale data is present during refetching, even if it errors out
    if (!isFetching && data) {
      return {
          status: UserStatus.AUTHENTICATED,
          data
      };
    }
    return {
      status: UserStatus.UNKNOWN
    };
  },
  async (get, set, update) => {
    const { refetch } = get(getUserAtom);
    if (update?.action === UserAtomActions.Refresh) {
      await refetch();
    }
  }
);

Feeling pleased with my code gymnastics, I started replacing the useUser context hook with useAtom and this atom. Now the entire app started re-rendering when a refetch was triggered.

Embracing the atomic way of thinking

I realized that the reason why the entire app was re-rendering was because the userAtom was not truly "atomic". It contained many pieces of state and each of them should be in a separate atom. If not, any change to the upstream getUserAtom would cause all components subscribed to userAtom to re-render.

So I split the userAtom into the following atoms.

// The logic for determining the status was slightly changed so that the state
// doesn't switch back and forth as a refetch is running
const userStatusAtom = atom(
  (get) => {
    const { data, error } = get(getUserAtom);
    if (error?.message === 'Authentication required') {
      return UserStatus.UNAUTHENTICATED;
    }
    if (!data) {
      return UserStatus.UNKNOWN;
    }
    if (data) {
      return UserStatus.AUTHENTICATED;
    }
  }
);

const userIdAtom = atom(
  (get) => {
    const { data } = get(getUserAtom);
    return data?.id;
  }
);

const usernameAtom = atom(
  (get) => {
    const { data } = get(getUserAtom);
    return data?.username;
  }
);

const userAvatarAtom = atom(
  (get) => {
    const { data } = get(getUserAtom);
    return {
      character: data?.avatar_character,
      background: data?.avatar_background
    };
  }
);

While fixing the re-renders, I discovered a few caveats about how the atom hooks behaved. If a derived atom returns a primitive value like a string, the hooks will only trigger a re-render if that value changes. However, if a derived atom returns an object, the hooks will trigger re-renders when any upstream atoms change. So in the case of the user atoms above, a hook which is tied to the usernameAtom will only trigger a re-render when the username changes. But if it was tied to the userAvatarAtom, it will trigger a re-render every time the getUserAtom changes, even if the avatar_character and avatar_background has not changed. So for the best performance, we need to be extremely "atomic" and avoid object atoms.

Atomize all things

Because I like making things harder for myself, I decided to go further and refactor the entire app to use jotai-tanstack-query atoms. It took me a while but it also gave me a chance to see some of Jotai's merits.

Using atoms allowed me to move all the data-processing logic out of the components. This wasn't possible before since Tanstack query's hooks returned the data within the component. Now all of that was happening within the atoms.

The atoms also reduced the amount of unnecessary computations. The best example of this was the Entries screen. Before switching to atoms, it had many computations within the render method that ran every time a hook within the component updated.

// Simplified excerpt of the Entries screen
// https://github.com/bashlk/timo/blob/main/packages/tanstack-query/src/routes/Entries/Entries.jsx
import { useQuery } from '@tanstack/react-query';
import { listEntries } from '@timo/common/api';

const Entries = ({ history }) => {
  const { data: entries } = useQuery({
    queryKey: ['entries'],
    queryFn: listEntries
  });

  // Other mutation and useState hooks

  return (
    <div className={styles['body']}>
      <div className={styles['total-row']}>
        <h2 className={styles['total-label']}>Total</h2>
        <div>{formatDuration(getTotalDuration(entries))}</div> {/* Computation 1 */}
      </div>
      {Object.entries(getEntriesGroupedByDate(entries)).map(([date, dayEntries]) => ( // Computation 2
        <div className={styles['day']} key={date}>
          <div className={styles['day-header']}>
            <h2 className={styles['day-name']}>{date}</h2>
            <div>{formatDuration(getTotalDuration(dayEntries))}</div> {/* Computation 3 */}
          </div>
          {dayEntries.toReversed().map((entry) => ( // Computation 4
            <Entry {...entry} />
          ))}
        </div>
      ))}
    </div>
  )
}

These computations could be avoided by using React's useMemo hook but that makes the code harder to read. Jotai's atoms were a much nicer way to get these computations to run only when the data changed.

// Simplified excerpt of the time entry atoms
// https://github.com/bashlk/timo/blob/main/packages/jotai/src/atoms/entryAtoms.js
import { atom } from 'jotai';
import { atomWithQuery } from 'jotai-tanstack-query';

const entriesAtom = atomWithQuery(() => ({
  queryKey: ['entries'],
  queryFn: listEntries,
}));

// This only recomputes when entriesAtom changes
const entriesGroupedByDateAtom = atom(
  (get) => {
    const { data } = get(entriesAtom);
    if (data) {
      return getEntriesGroupedByDate(data);
    }
    return {};
  }
);

// This only recomputes when entriesGroupedByDateAtom change
const entriesDurationsGroupedByDateAtom = atom(
  (get) => {
    const entries = get(entriesGroupedByDateAtom);
    if (entries) {
      return Object.entries(entries).reduce((grouped, [date, dayEntries]) => {
        grouped[date] = getTotalDuration(dayEntries);
        return grouped;
      }, {});
    }
    return {};
  }
);

// Simplified excerpt of the reworked Entries screen
// https://github.com/bashlk/timo/tree/main/packages/jotai/src/routes/Entries
import { useAtomValue } from 'jotai';
import {
  entriesGroupedByDateAtom,
  entriesDurationsGroupedByDateAtom,
  entriesDurationAtom,
} from '../../atoms/entryAtoms';

const Entries = ({ history }) => {
  // These atoms are not recomputed during every render
  const groupedEntries = useAtomValue(entriesGroupedByDateAtom);
  const entriesDuration = useAtomValue(entriesDurationAtom);
  const entriesDurationsGroupedByDate = useAtomValue(entriesDurationsGroupedByDateAtom);

  return (
    <div className={styles['body']}>
      <div className={styles['total-row']}>
        <h2 className={styles['total-label']}>Total</h2>
        <div>{entriesDuration}</div>
      </div>
      {Object.entries(groupedEntries).map(([date, dayEntries]) => (
        <div className={styles['day']} key={date}>
          <div className={styles['day-header']}>
            <h2 className={styles['day-name']}>{date}</h2>
            <div>{entriesDurationsGroupedByDate[date]}</div>
          </div>
          {dayEntries.map((entry) => (
            <Entry {...entry} />
          ))}
        </div>
      ))}
    </div>
  )
}

Summary

I am conflicted about Jotai and atomic state management. On one hand, it provides a very effective way to move stateful logic outside of React components. But on the other, it is very easy to blow your leg off. A bad decision about how the atoms should be separated could lead to re-renders all across the app. The way how Jotai's atoms and hooks behave is not very obvious. Jotai's utilities and extensions are not described in full depth in the docs so some exploration is needed to figure out how to use them for more advanced scenarios.

So I can't recommend Jotai unless you really love the idea of working with atoms. It took me a while to get used to them since they need a different way of thinking than the common flux-like state management. In the next monthly post, I will be taking a look at another unconventional form of state management by trying out Valtio. If you are new here, you can sign up for the monthly posts so that you don't miss the next one.