Jotai: The ultimate React State Management

Jotai is a relatively new state management library for React. It's simple, but make no mistakes, it's a robust library.

Jotai is based on the new Recoil pattern and library by Facebook. More than half a decade ago, Facebook created a pattern and library for state management in React called Flux .

On this pattern, some non-Facebook devs created a new library of their own that was more robust, simpler to use, and took the React world by storm. This library is Redux . Now Facebook has Recoil, which has different ideologies than Flux.

The same is the case with Jotai and Redux. Let's explore some of these.

Why Jotai?

Minimalistic API - Jotai has a simple API design and is a joy to work with.

Tiny bundle size - Jotai's footprint is really small and won't be a performance block in your site/app

Loaded to the brim - Jotai comes with a lot of

Performant - Jotai is BLAZING FAST. It's runtime performance is insane!

TYPESCRIPT!! 🥳🥳 - First class TypeScript Support!! Comes with Typings pre-installed, and TypeScript authoring experience is beyond heavenly.

Ideological differences from Redux

Jotai is very different from Redux and React Context API in almost every way. But there's one central concept that is the catch-all—the one that you need to internalize.

Redux stores are monolithic, but Jotai is atomic.

This means, in Redux, it's a pattern to store all the needed global state in the app in one single big object. In Jotai, it’s the opposite. You break your state into atoms i.e. - one store for one single store, or for closely related state.

Getting started with Jotai

Installing Jotai

# pnpm
pnpm add jotai

# npm
npm install jotai

# Or if you're a yarn person
yarn add jotai

Setting it up in the app

Jotai requires a provider to be present in the parent of the current component where it is being used. The simplest way to do would be to wrap the whole app in the provider as shown below 👇

// index.jsx (or index.tsx)
import React from 'react';
import ReactDOM from 'react-dom';
import { App } from './App';

// Jotai provider
import { Provider } from 'jotai';

ReactDOM.render(
  <React.StrictMode>
    <Provider>
      <App />
    </Provider>
  </React.StrictMode>,
  document.getElementById('root'),
);

Now you can use jotai anywhere in your app!

Basic syntax

Now that our basic setup is done let's have a look at the syntax!

Creating your first atom

Atoms are the building blocks of universe and clump together into molecules--

No, not that atom 😅.

Jotai atoms are small isolated pieces of state. Ideally, one atom contains very small data (although it's just a convention. You could still put all your state in one atom, though that would be very slow performance-wise).

So here's how you create your first atom

import { atom } from 'jotai';
const themeAtom = atom('light');

And that's it! You have your very first piece of state in!!

Notice I suffixed my atom name with Atom, as in themeAtom. It's not a rule or an official convention., I simply choose to name my atoms like this for clarity in a big project. You can name it just theme rather than themeAtom 🙂

Now, how do you use it? Well, using it is a cross between the useState and useContext hooks.

import { useAtom } from 'jotai';

export const ThemeSwitcher = () => {
  const [theme, setTheme] = useAtom(themeAtom);

  return <main>{theme}</main>;
};

See? Exactly like useState, but the only difference is that the atom we created is to be passed to useState. useAtom returns an array of size 2, where the 1st element is a value and the 2nd element is a function, to set the value of the atom. This makes all the components relying on this atom update and re-render.

So, if we put it all together, the complete code would look like this 👇

import { atom, useAtom } from 'jotai';

const themeAtom = atom('light');

export const ThemeSwitcher = () => {
  const [theme, setTheme] = useAtom(themeAtom);

  return <main>{theme}</main>;
};

And notice the setTheme isn't used yet. Let's change that 👇

import { atom, useAtom } from 'jotai';

const themeAtom = atom('light');

export const ThemeSwitcher = () => {
  const [theme, setTheme] = useAtom(themeAtom);

  return (
    <main>
      <p>Theme is {theme}</p>

      <button onClick={() => setTheme(theme === 'light' ? 'dark' : 'light')}>Toggle Theme</button>
    </main>
  );
};

Check it out @ CodeSandbox

This, my friend, is just the beginning. Jotai can do so much more!

But this alone doesn't give much perspective. What's special about a button that toggles a value 🙄? And I agree. This example is pretty boring. Let's use Jotai to make an actual theme switcher.

Jotai In Practice: Theme Switcher hook

Nowadays, theme switching is needed in every single app, website, and heck, even blog sites (Especially blogs). And it can be quite daunting to make a theme switcher. First, you have to set up your CSS Variables. Then you have to start with a theme, post which you have to make a button that switches the theme. Then you have to make sure to remember the preference using localstorage API. But that brings you full circle to picking up the right value when the page loads, and also not messing with SSR and prerendering and...

Yeah, it's pretty complicated. A problem any developer would dread before attempting (I did 🥶).

So, that's the best kind of thing to make, and let's use Jotai to do it. You'll be astonished at how simple Jotai can make it.

So, here are our goals:

Works on server side (As in not referring to document or window without protection).

Picks up locally stored value in localstorage.

If no local value, tries to get the device preference, whether the device theme is light or dark.

Current theme should be available as a state that re-renders components it is being used in.

Changing the state should update localstorage accordingly.

So now that our list is complete, let's look at the code 👇

import { atom, useAtom } from 'jotai';
import { useEffect } from 'react';

const browser = typeof window !== 'undefined';

const localValue = browser ? localStorage.getItem('theme') : 'light';
const systemTheme =
  browser && matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light';

// The atom to hold the value goes here
const themeAtom = atom(localValue || systemTheme);

/** Sitewide theme */
export function useTheme() {
  const [theme, setTheme] = useAtom(themeAtom);

  useEffect(() => {
    if (!browser) return;

    localStorage.setItem('theme', theme);

    document.body.classList.remove('light', 'dark');
    document.body.classList.add(theme);
  }, [theme]);

  return [theme, setTheme];
}

A lot going on here. So here's the breakdown.

We check whether the current code is running in the browser or not. If we're running the code in SSR or prerendering, this value will be false.

We get the value stored in localstorage. If localstorage has the theme in it, then we'll consider it the highest priority, as it would be the one chosen by the user. Also, because there's no localstorage in Node, we have to fallback to the default value of light if it's running in the SSR mode.

We also retrieve the device preference using prefers-color-scheme: dark, in case the localstorage value doesn't exist. Again, this falls back to value light if device preference is dark or code is running in SSR.

Finally, create the atom. This will be our main store where we actually store the current theme, usable and changeable as state. Notice the value we give it: localValue || systemTheme. Here's what can happen with these values:

If running in SSR/Prerendering mode, localValue = 'light' and systemTheme = 'light', localValue || systemTheme will turn out to be light. So, important point here: Your app in SSR will be themed with light theme, so if you prerender your app, it will end up with light theme, in terms of plain HTML. As the JavaScript loads, it will sync to the most relevant theme possible.

Why didn't I just put the localValue and systemTheme variables inside the hook? The reason: If I put them in the hook, everytime the hook is initialized in any component, or a component re-renders, this hook will run again, and will fetch these values again from localstorage and media queries. These are pretty fast, but localstorage is blocking, and when used a lot, can introduce jank. So we initialize these 2 vars once in the lifetime of the app, because we need these only to get the initial value.

Finally, let's begin our hook:

Let's make this atom a local state using useAtom: const [theme, setTheme] = useAtom(themeAtom);. These will be our theme in the form of state. Themes can be modified using setTheme.

Next thing, we got the most important part of our hook that will actually make the current theme known to our CSS.

useEffect(() => {
  if (!browser) return;

  localStorage.setItem('theme', theme);

  document.body.classList.remove('light', 'dark');
  document.body.classList.add(theme);
}, [theme]);

It's useEffect that runs whenever theme changes, as you can see in the array in the 2nd argument. When this runs, it checks if the code is running in the browser. If it isn't, it simply stops further execution by doing a return.

If it is successful, it goes on and removes all the classes corresponding to put themes on <body>, then it adds the class corresponding to the latest value of theme variable.

Finally, we return the [theme, setTheme] pair as it is, so we can use it just like we use useState. You could also return these as objects { theme, setTheme } giving them explicit naming.

This is it for this hook.

And I got my TypeScript kin covered too 😉👇

import { atom, useAtom } from 'jotai';
import { useEffect } from 'react';

export type Theme = 'light' | 'dark';

const browser = typeof window !== 'undefined';

const localValue = (browser ? localStorage.getItem('theme') : 'light') as Theme;
const systemTheme: Theme =
  browser && matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light';

// The atom to hold the value goes here
const themeAtom = atom<Theme>(localValue || systemTheme);

/** Sitewide theme */
export function useTheme() {
  const [theme, setTheme] = useAtom(themeAtom);

  useEffect(() => {
    if (!browser) return;

    localStorage.setItem('theme', theme);

    document.body.classList.remove('light', 'dark');
    document.body.classList.add(theme);
  }, [theme]);

  return [theme, setTheme] as const;
}

So this is the final code we got for switching themes robustly. This hook is simple and understandable (I hope 😅), due to the simplicity of Jotai.

But here's the thing. We haven't saved much code using Jotai compared to using the Context API. That code would be almost this simple, just have a little more boilerplate. So there really isn't much difference here.

But, here comes a twist: We can get rid of even more code by using something that Jotai provides: atomWithStorage

We can move the logic of syncing to localstorage completely, both from inside the hook as well as outside.

Rewriting the hook with atomWithStorage

atomWithStorage is a special kind of atom that automatically syncs the value provided to it with localstorage or sessionStorage (Or AsyncStorage, if used with React Native), and picks the value upon the first load automatically! It's available in the jotai/utils module, and adds some bytes other than the 2.4KB of Jotai Core.

So here's how we would rewrite it:

import { useAtom } from 'jotai';
import { atomWithStorage } from 'jotai/utils';
import { useEffect } from 'react';

const browser = typeof window !== 'undefined';

// The atom to hold the value goes here
const themeAtom = atomWithStorage(
  'theme',
  browser && matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light',
);

/** Sitewide theme */
export function useTheme() {
  const [theme, setTheme] = useAtom(themeAtom);

  useEffect(() => {
    if (!browser) return;

    document.body.classList.remove('light', 'dark');
    document.body.classList.add(theme);
  }, [theme]);

  return [theme, setTheme];
}

As you can see, we completely got rid of localstorage from the code, and we have a new thing atomWithStorage. First argument is the key to store it in localstorage. As in, if you specified theme as value here, you would retrieve it from localstorage using localstorage.getItem('theme').

As you can see, the code itself isn't much smaller in terms on lines of code. It's just 20% smaller, which isn't a big number in case of this already-small file. The main part here is that we got to hide the complexity away thanks to atomWithStorage. Now we don't have to keep the local value storage in mind, just have to focus on our main logic and remember that this value is synchronized locally, and that's it.

And using this hook ultimately is super simple,

import { useTheme } from './use-theme';

export const ThemeSwitcher = () => {
  const [theme, setTheme] = useTheme();

  return (
    <main>
      <p>Theme is {theme}</p>

      <button onClick={() => setTheme(theme === 'light' ? 'dark' : 'light')}>Toggle Theme</button>
    </main>
  );
};

And it just works!! 🪄

More about Jotai

So this was the basic intro to Jotai. I added in the atomWithStorage utility function to show how powerful and simple it can make your code. I will touch on these utilities later on. For now, let's explore more about the basic atom and useAtom, and how they give you superpowers.

Derived atoms

Sometimes, you want to make an atom rely on another atom(s), meaning -you want to compose multiple atoms together into one big computed atom. That is extremely simple with Jotai.

Readonly atoms

Readonly atoms are derived atoms that rely on other atoms, and we can't change their values directly.

For example, these atoms' usage would be as follows 👇

const [derivedValue] = useAtom(derivedAtom);

There's no setDerivedValue here, no setter function. We can only read this atom. Changing the atoms it is derived from will automatically update this value.

But enough talk! Now let's see how to create these derived atoms.

You have seen this atom until now 👇

const store = atom('someValue');

But guess what? Atoms can take a function as a parameter 👇

const store = atom((get) => get(someAtomDefinedSomewhere));

Here, instead of a raw value, we pass a callback. This callback has a parameter get, which allows you to access the raw, live values of other atoms. With this raw value, you can do anything. Multiply it, concatenate it, map it, reduce it, whatever. The sky’s the limit.

And you can do a lot more with this. For example, one simple example would be to have a list of all the keys of an object that match a specific criteria to be in an array.

Here's the object

export const appsStateStore = atom({
  finder: false,
  launchpad: false,
  safari: false,
  messages: false,
  mail: true,
  maps: true,
  photos: false,
  facetime: true,
  calendar: false,
});

Define the atom that will hold the open apps in an array 👇

const openAppsStore = atom((get) => {
  const apps = get(openAppsStore); // Gives the raw value { finder: false, launchpad: false, ...

  // Filter out the values who are marked as false
  const openAppsList = Object.keys(apps).filter((appName) => apps[appName]);

  return openAppsList;
});

And this is it!! As you tweak the values in the appStateStore, setting them to true and false, the openAppsStore will reflect the changes and the components using this store will also be updated with new values.

You can also compose together many different atoms together 👇

const xCoordinateAtom = atom(0);
const yCoordinateAtom = atom(0);

// Compose 'em all
const distanceFromOriginAtom = atom((get) =>
  Math.sqrt(get(xCoordinateAtom) ** 2 + get(yCoordinateAtom) ** 2),
);

You can tweak the xCoordinateAtom atom and yCoordinateAtom, and the distanceFromOriginAtom will update with the new values!!)

It's a mathematical formula to calculate the distance of a point from the origin (0, 0). If you didn't get it, no worries, I just want to get the point across that you can compose together different atoms seamlessly. That's it! 🙂

Readable & Writable atoms

These are atoms that are derived from other atoms, but can also be modified on their own by the user.

const readWriteAtom = atom(
  (get) => get(priceAtom) * 2,
  (get, set, newPrice) => {
    set(priceAtom, newPrice / 2);
    // you can set as many atoms as you want at the same time
  },
);

This atom, when you set its value, triggers the custom write function we provide, and can modify the atoms it relies on. It's basically two-way data binding. You change priceAtom, this readWriteAtom gets updated. You update readWriteAtom, priceAtom gets updated. Mindblowing, right 🤯🤯?!?

Beware though: As magical as this seems, it's two-way data binding. There have been controversies in the past about it, and rightfully so, as debugging and keeping the flow of data sane becomes extremely hard with these. That's why React itself has only one-way data binding. So use this atom carefully.

Async Atoms

From this point, we enter a very dangerous territory: Async rendering, aka React Suspense.

Sometimes your atoms have to be asynchronous, that is, rather than getting values instantly, they pull from a remote source using fetch, which is when you have to suspend the rendering and wait for the data to come back.

Here's a little code demonstration of using async atom 👇

const fetchCountAtom = atom(
  (get) => get(countAtom),
  async (_get, set, url) => {
    const response = await fetch(url);
    set(countAtom, (await response.json()).count);
  },
);

function Controls() {
  const [count, compute] = useAtom(fetchCountAtom);
  return <button onClick={() => compute('http://count.host.com')}>compute</button>;
}

But this above won't work if you don't wrap Controls in a Suspense 👇

<Suspense fallback={<span />}>
  <Controls />
</Suspense>

Async Atoms are extremely useful in building real world apps, because these apps are mostly CRUD apps with data fetching added in.

The Best of Utils

If you loved atomWithStorage and your head is spinning with all the possibilities it could unlock, I got many more awesome Jotai utils for you.

atomWithStorage

I covered this one in the very beginning of the article when I refactored the useTheme hook to use this special atom. It accepts a key (The name by which it is stored in localstorage), and the initial value. Then you change this atom, and its value will be persisted locally and picked up after the page reloads.

import { atomWithStorage } from 'jotai/utils';

const darkModeAtom = atomWithStorage('darkMode', false);

This atom is also SSR friendly, so you can SSR your app away with absolutely no issues.

This atom can store value in sessionStorage too, so the atom's value will be persisted until the browser is closed. Handy if you're building a banking web app, where having short sessions is preferable.

It also works with React Native, so it's pretty much universal 🤩

atomWithReset

Sometimes you need to reset your state to what it was originally. Traditionally, the way to do that has been to store the initial value in a variable, create a state with that variable as the value, and when needed, setState back to that initial value. The code would look like this 👇

import { atom, useAtom } from 'jotai';

const initialValue = 'light';

const themeAtom = atom(initialValue);

function ThemeSwitcher() {
  const [theme, setTheme] = useAtom(themeAtom);

  const toggleTheme = () => setTheme(theme === 'light' ? 'dark' : 'light');
  const resetTheme = () => setTheme(initialValue);

  return (
    <>
      <button onClick={toggleTheme}>Toggle theme</button>

      <button onClick={resetTheme}>Reset theme</button>
    </>
  );
}

This is fairly easy, but here's a more Jotai-ish way of doing the same thing 👇

import { useAtom } from 'jotai';
import { atomWithReset, useResetAtom } from 'jotai/utils';

const themeAtom = atomWithReset('light');

function ThemeSwitcher() {
  const [theme, setTheme] = useAtom(themeAtom);
  const reset = useResetAtom(themeAtom);

  const toggleTheme = () => setTheme(theme === 'light' ? 'dark' : 'light');

  return (
    <>
      <button onClick={toggleTheme}>Toggle theme</button>

      <button onClick={reset}>Reset theme</button>
    </>
  );
}

As you can see, we simplified our component a little. Not much, in this case, as it's a very simple example. But I personally have used this reset atom in my apps with full blown complex logic based components, and it just makes the code much more sane, idiomatic, and bug-free.

selectAtom

If there was a coolness counter for libraries and frameworks, Jotai alone would've broken it with this little util.

Let’s consider you have a big object.

const defaultPerson = {
  name: {
    first: 'Jane',
    last: 'Doe',
  },
  birth: {
    year: 2000,
    month: 'Jan',
    day: 1,
    time: {
      hour: 1,
      minute: 1,
    },
  },
};

// Original atom.
const personAtom = atom(defaultPerson);

And say, a lot of components rely on this specific atom, but need only parts of this.

The thing is when you update this atom, all the components relying on this atom will rerender. Even if you change just the birth.time.minute, the whole thing is going to count as an update and all the components will re-render. This is how React works, unfortunately.

But worry not, for Jotai has a solution for this too! selectAtom allows you to create a derived atom with only a subpath of the whole object.

const firstNameAtom = selectAtom(personAtom, (person) => person.name.first);

firstNameAtom is a read-only derived atom that only triggers when the person.name.first property changes, and it holds the value of person.name.first.

You can update the birth.time.hour field(By updating the whole atom with new values), and the component relying on firstNameAtom will remain unchanged. Amazing, right?

Applying on Objects

There arises a problem: If you listen to a field that is an object, person.birth, this atom isn't gonna be very efficient. Jotai uses the equality check(===) to check if the atom's part is changed or not and should be re-rendered. The thing is, no 2 objects are ever the same. The === checks objects by reference, not values. So basically, this atom is pretty useless in that scenario. But not quite!

You can provide a 3rd argument to this selectAtom, which is your own version of an equality check. You can write your custom function to check the objects.

const birthAtom = selectAtom(personAtom, (person) => person.birth, deepEqual);

OFC, writing your own deepEqual is hard, so it's recommended to go with lodash-es's isEqual function.

import { isEqual } from 'lodash-es';

const birthAtom = selectAtom(personAtom, (person) => person.birth, isEqual);

If seeing lodash gives you anxiety about bundle size, I assure you, isEqual of lodash-es is tree-shakeable, just 4.4KB minified, and even smaller in gzip/brotli. So no worries 😁

This can take the performance of your app from zero to hero. Literally!

freezeAtom

import { atom } from 'jotai';
import { freezeAtom } from 'jotai/utils';

const objAtom = freezeAtom(atom({ count: 0 }));

freezeAtom takes an existing atom and returns a new derived atom. The returned atom is "frozen" which means when you use the atom with useAtom in components or get in other atoms, the atom value will be deeply frozen with Object.freeze. It would be useful to find bugs where you accidentally tried to mutate objects which can lead to unexpected behavior.

This atom is mostly for debuggability, for when you mutate an object state(which you aren't supposed to do in React, but hey, we're all humans). This is such a common case, that I'm really glad Jotai folks are providing such high quality debugging tools.

waitForAll

Remember the section above about Async atoms? This util is for that, and quite a handy one it is.

const dogsAtom = atom(async (get) => {
  const response = await fetch('/dogs');
  return await response.json();
});

const catsAtom = atom(async (get) => {
  const response = await fetch('/cats');
  return await response.json();
});

const App = () => {
  const [dogs] = useAtom(dogsAtom);
  const [cats] = useAtom(catsAtom);
  // ...
};

So you have these 2 async atoms, and you're using them in the app. All fine. But there's a little problem here: The component will wait for the first atom dogsAtom to go fetch data, return, then it will move to the next atom catsAtom. We don't want this. Both these atoms are independent of each other, we should rather fetch them in parallel(Or concurrently if you are a hardcore JavaScripter 😉)

So, we basically want to do something like an await Promise.all(...) on these atoms. The way to do that is using the waitForAll util.

After using, our code becomes 👇

const dogsAtom = atom(async (get) => {
  const response = await fetch('/dogs');
  return await response.json();
});

const catsAtom = atom(async (get) => {
  const response = await fetch('/cats');
  return await response.json();
});

const App = () => {
  const [[dogs, cats]] = useAtom(waitForAll([dogsAtom, catsAtom]));
  // ...
};

Now it waits for both of them to resolve, and then returns an array of the data returned by both. Kind of like an await Promise.all statement.

Literally, at this point, React should absorb Jotai into itself, it's way too good!!

And these are only half of all utils provided by Jotai. There are so many, I'd could write a whole book about it . Head over to Jotai Documentation to learn about em.

Jotai is good with its relatives 🤝

Jotai isn't anything like other libraries. which are like: "You shall only have me in thy package.json!!!"

No, Jotai doesn't work like that! Jotai itself is a great state management library, but it also allows you to seamlessly integrate with other state management libraries.

Here are all the official integrations that come with Jotai:

Immer

Optics

React Query

XState

Valtio

Zustand

Redux

URQL

Now, at this point, the blog post is already way too lengthy to cover the above integrations, but I want to cover Immer. Why? Because of the biggest pain point with React state: Immutability.

Immutability is great, and it makes wrapping your head around React State easy, but it can make things very hard when your state is an object. Then you have to do the whole song and dance of spreading the object and merging with the properties you want to update.

function UpdateUser() {
  const [user, setUser] = useState({
    id: 23,
    name: 'Luke Skywalker',
    dob: new Date('25 December, 19 BBY'),
  });

  // Update the dob
  const updateDob = () => setUser({ ...user, dob: new Date('25 November, 200ABY') });

  return <button onClick={updateDob}>Update DOB</button>;
}

As you can see in the updateDob method, we have to spread the original object, and pass the field we want to update. This is OK. But what if the object is many levels deep and we want to update an object very deep.

It becomes so convoluted that I personally never even tried it. I just re-architected my state to be shallower in some way and then update that. I am more of a Svelte person than a React person, and in Svelte, you can simply mutate the state and it just works.

user.dob = new Date('25 November, 200ABY');

And it goes extremely deep too!

state.depth1.depth2.depth3.depth4 = 'something';

So all the song and dance required in React always felt wrong to me.

But this is where Immer comes in. Immer allows you to directly mutate the state, and it just works. Take a look for yourself

import { atomWithImmer } from 'jotai/immer';

const userAtom = atomWithImmer({
  id: 23,
  name: 'Luke Skywalker',
  dob: new Date('25 December, 19 BBY'),
});

function UpdateUser() {
  const [user, setUser] = useAtom(userAtom);

  // Update the dob
  const updateDob = () =>
    setUser((user) => {
      user.dob = new Date('25 November, 200ABY');
      return user;
    });

  return <button onClick={updateDob}>Update DOB</button>;
}

Here, the setUser works differently. It's a callback that passes you the current value of the state. This value is a copy of the original value. You can mutate this copy as much as you want inside the callback, and finally just return it, Jotai and Immer will automatically reconcile the changes without any of the bugs that come with mutating. Freaking awesome!

Conclusion

Alright!! This was a lengthy article! Congratulate yourself for making it to the end (skimming through counts too 😉).

This article is just a little glimpse of Jotai. There's so much more at Jotai docs. You should definitely check them out.

Peace out✌️