State Management in React (2/3) - Jotai

Jun-09-2024

React State Management

This is the follow-up of a previous post related to React global state management. In the previous post we looked at React Context. In this post we will look at Jotai, which is a library to manage global state.

Application Goal

We will aim to work on the same example project than the previous case, i.e. implementing a simple version of a Kanban table.

Jotai: A simple view

Overall principle

Jotai doesn’t work with a global store containing all the states, but it works with declaring an atom with each piece of state. Basically each value you want to keep track of is its own store, a bit like a supercharged useState you could share easily among the component tree.

Some key functions

First you need to define an atom, which will contain the piece of state. Other files will need to import this atom, so it’s better to put this declaration in its own file.

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

Indeed very similar to a useState declaration.

When using this value, we can just import this atom and use the hook useAtomValue.

For example,

const taskList = useAtomValue(taskListAtom);

This will cause the component to re-render when the atom value changes.

In the case we want to modify the value, we need to use the hook useSetAtom which will return a setter function.

const setTaskList = useSetAtom(taskListAtom);

// ...

// when need to add a new task
setTaskList((taskList) => {
    return [...taskList, newTask];
})

The setter callback will receive the current value in parameter and should return the new atom value.

All in all it is pretty simple and not verbose - I find this solution really elegant.

React Suspense

Since React 18, React is offering the possibility of Suspense components, i.e. offering the possibility of rendering a fallback component while some process is underway, and automatically displaying the computed value when the process is finished. However, this worked only for react server components (that you cannot easily re-create without a framework like Next or Remix) or for lazy loading some components. But Jotai is actually compatible with this, so that you can use React Suspense in conjunction with Jotai.

Application to the Kanban project

Here is the basic way to use jotai inside the Kanban project:

Store declaration

Inside the store we will create the atom and export it.

export const TaskListAtom = atom<Task[]>([]);

Because of the way we interact with the tasks (i.e. just bump up or bump down the task priority), we will add a custom React hook which will simplify the code for us. The hook will return an array with two modifiers, one to decrease the task priority and the other one to increase it.

// Importing the atom defined above
import { useTaskListCommands } from './TaskStore.tsx';

export const useTaskListCommands = () => {
  // Jotai hook "useSetAtom"
  const setTaskListAtom = useSetAtom(TaskListAtom);

  // function which increases the priority of a task
  const upgradeTask = (taskName: string) => {
    setTaskListAtom((currentTaskList) => {
      // This hook takes the current task list as parameter
      // and should return the new list
      const oldTask = currentTaskList.find((task) => task.name === taskName);
      if (!oldTask) {
        return currentTaskList;
      }
      const levelIndex = TaskStateList.findIndex(
        (taskState) => taskState === oldTask.state
      );
      // Looking up the next state
      const newLevel =
        TaskStateList[Math.min(levelIndex + 1, TaskStateList.length - 1)];
      oldTask.state = newLevel;
     // using the spread operator to return a new object with the same
     // elements
     return [...currentTaskList];
    });
  };
  const downgradeTask = (taskName: string) => {
    // skipped for brevity
  };

  return [downgradeTask, upgradeTask]; 

Once this is declared, we import the value or the functions only in the components where we would use them:

  • First in the parent component, to access the whole list in a read-only way
import { TaskListAtom } from './TaskStore.tsx';

// inside the component rendering code
const currentTaskList = useAtomValue(TaskListAtom);
  • Then for each task component, using the custom hook to access the modifiers:

// inside the component rendering code
const [downgradeTask, upgradeTask] = useTaskListCommands();

  return (
    { /* details omitted */ }
        <button onClick={() => downgradeTask(task.name)}>-1</button>
        <button onClick={() => upgradeTask(task.name)}>+1</button>
    { /* details omitted */ }
);

You can see the whole code at https://stackblitz.com/edit/vitejs-vite-6lrazu?file=src%2FTaskBox.tsx

Conclusion

jotai manages to provide a global state with very little boilerplate code. Importing a third-party library has some downside as well - and for basic uses you might be happy with the basic React Context version. I don’t really believe in recommendations, as your mileage may vary. What is more important is understanding the tools so that you can make an informed choice if/when you have to decide your stack. Another worthy goal is just to experiment: using a different tool forces you to look at a problem you are familiar with, from an unfamiliar angle. And challenging yourself this way is always good to deepen your understanding and to flex your mental models.