jotai-optimistic

A highly opinionated approach to optimistic updates with Jotai and Immer.

Mutate your state optimistically, run your network request, and don't worry about rolling back the optimistic update if the network request fails.

Motivation

Optimistic updates are hard.

In order to implement them properly, you need to:

1. Make some set of changes to your state
2. Kick off some asynchronous action to sync that state with the server
3. If there's an error, roll back just the changes you made in step 1

Step 3 is of course the hard part, especially if your action is modifying some deeply nested attribute of a larger state object. If you do some naive implementation like:

1. Take a snapshot of your state
2. Make some set of changes to your state
3. Kick off some asynchronous action to sync the state
4. If there's an error, restore the snapshot you took in step 1

You will have a bug that involves resetting any other changes to the state that happened in the interim. No good!

Solution

The solution is a hook, called useAtomImmerSaga. Here's what it looks like to use it in a section of code responsible for updating the value of a toggle representing whether a particular relationship is directed or has no direction.

1export const useRelationshipKindHasDirection = (
2  id: IDTypes["relationshipKind"]
3) => {
4  const [relationshipKind, runSaga] = useAtomImmerSaga(
5    relationshipKindByIdAtomFamily(id)
6  );
7
8  const setHasDirection = (hasDirection: boolean) =>
9    runSaga((saga) =>
10      saga
11        .update((draft) => {
12          draft.has_direction = hasDirection;
13        })
14        .effect(async (_nextState, _relationshipKind) => {
15          await trpc.updateRelationshipKind.mutate({
16            id: id,
17            patch: { has_direction: hasDirection },
18          });
19        })
20        .onError((draft, error) => {
21          draft.error = error.toString();
22        })
23    );
24
25  return [relationshipKind.has_direction, setHasDirection] as const;
26};

And this is what it looks like to build a component using that hook:

1const EditableRelationshipHasDirection = ({ id }: {
2  id: IDTypes["relationshipKind"];
3}) => {
4  const [hasDirection, setHasDirection] = useRelationshipKindHasDirection(id);
5
6  return (
7    <button
8    type="button"
9    onClick={() => setHasDirection(!hasDirection)}
10    />
11  );
12};

I think it's pretty great! You get a hook that abstracts the network call, the application of the optimistic update, and its rollback in the event of a network failure.

The key bit is the typed saga, which has .update, .effect, and `.postEffect`` methods.

The .update method is applied immediately - that's the optimistic state update, which, thanks to immer, you can just apply via easy imperative object mutation.

The .effect method contains the network call or other asynchronous side effect of the user action. If it throws, the changes applied during the .update method and only those changes will be rolled back. The full state will not be reset to what it was before the network mutation.

There is also a .postEffect method for applying some state update after the network call has succeeded. I was originally using it to plug in a server generated ID, but I have since switched to using client side generated branded IDs for my particular project. I'm going to keep it around for a while to make sure I don't need it for anything else.

Other Exports

There are a few other exports here that are useful.

createDerivedImmerAtom

This function can help create a derived, "immerified" atom from a larger atom that you can run optimistic updates with.

1import { createDerivedImmerAtom } from 'jotai-optimistic';
2
3const bigAtomWithNestedObjects = atomWithImmer({
4  bigListOfEntities: [
5    {
6      name: 'a',
7      count: 42
8    },
9    {
10      name: 'j',
11      count: 89
12    }
13  ],
14  anotherObject: {
15    nestedDate: new Date(),
16    nestedNumber: 1
17  }
18});
19
20const aAtomWithImmer = createDerivedImmerAtom(
21  bigAtomWithNestedObjects,
22  bawno => bawno.bigListOfEntities.find(entity => entity.find(name === 'a'))
23);
24
25const anotherObjectAtom = createDerivedImmerAtom(
26  bigAtomWithNestedObjects,
27  bawno => bawno.anotherObject
28);

aAtomWithImmer is now writeable, and nicely write-able with Immer style draft functions, and I haven't had to write any setter for it.

For example, I can do:

1const EditAnotherObjectName = () => {
2  const [anotherObject, setAnotherObject] = useAtom(anotherObjectAtom)
3
4  return (
5    <input value={anotherObject.name}
6      onChange={ev => {
7      setAnotherObject(draft => {
8        draft.name = ev.target.value
9      })
10      }
11    />
12  );
13}

useSetInitialAtomValueFromQuery

With derived atoms, the possibility of an atom being undefined can be really annoying, since you have to handle it for it with every single derivation. This is also true if an atom is async - every atom which reads from it has to be async as well.

I have found it easier to instead initialize the atom to an empty but not undefined object state, and then use useSetInitialAtomValueFromQuery to set the initial value of the atom from a query once it comes back.

Suppose trpc.getDocuments() returns { id: string, title: string, body: string }[]

I would:

1import { useAtomValue } from 'jotai';
2import { atomWithImmer } from 'jotai-immer';
3import { useSetInitialAtomValueFromQuery } from 'jotai-optimistic';
4
5const documentsAtom = atomWithImmer([]);
6
7const MyComponent = () => {
8  const documents = useAtomValue(documentsAtom);
9
10  const { data: documentsData, loading: documentsAreLoading } = trpc.useQuery.getDocuments()
11
12  useSetInitialAtomValueFromQuery(
13    documentsAtom, 
14    documentsData, 
15    documentsAreLoading
16  )
17
18
19  return documents.map(
20    // some list of documents
21  )
22}