Create your own useOptimistic hook

---

Recently, I came across this video by Web Dev Simplified, which introduced me to React’s new useOptimistic hook.

I began to wonder whether this is even needed. React has never been a library to provide more than bare-bones controls and hooks. The useCallback, useMemo, and useState hooks are most of the built-in hooks and they provide a solid base upon which you can build pretty much any logic inside your components or custom hooks.

So since when does React provide more high-level abstractions above hooks? At that point, they become even more opinionated than they already are. Not all use cases are the same - but this is a pretty generic hook, all things considered. Apparently, it is still in the experimental stage.

I decided to experiment with creating this hook on my own, to understand if it was overcoming some internal system limitation. I looked at various examples of existing optimistic hooks on NPM, but I didn’t really find one that falls in line with more modern coding styles, or they were completely outdated and abandoned.

Why create your own?

Ideally, React should generally keep its size down as much as possible, and not increase the bundle with more bloat that won’t be used by everyone. I believe this hook is a very nice idea, but it is not relevant for everyone, and I think if you find it does fit your needs, you don’t have to wait for official support - you can do it (or copy the final code here) in relatively short time, and feature-complete (and added).

In addition, I thought it would make a good thought experiment and some coding practice, it can’t hurt!

Looking at React’s original implementation

I decided to first take a look at React’s official implementation of this. Not the source code, but I did want to get a sense of whether I like the consensus they came up with, or whether I feel it has room for improvement.

Here is a basic usage example, verbatim from React’s docs:

const [optimisticState, addOptimistic] = useOptimistic(
  state,
  // updateFn
  (currentState, optimisticValue) => {
    // merge and return new state
    // with optimistic value
  },
)

Seems nice, but I think we can do better.

Off the top of my head, I don’t see a way to get the status of the optimistic value. Can I know when it’s been finalized without having to maintain a separate state?

Also, there doesn’t seem to be a way to roll back to the original value if an error has occurred. Or at least, the docs don’t yet specify any. We can fix that, and provide a more “holistic” approach, which in my opinion doesn’t leave the scope of this hook. Say I push an optimistic update, and it fails, how can I roll it back?

We’ll try to make sense of where it should go as we go along.

Creating a hook

Let’s start with the signature, and make sure we have our types going for us.

type StateUpdater<T, A = T> = (currentState: T, newState: A) => T

export function useOptimistic<T, A>(
  initial: T,
  merge?: StateUpdater<T, A>,
): readonly [T, React.Dispatch<React.SetStateAction<T>>] {
  // ...
}

We start with the basics. An initial value, of type T which is generic, and an optional merge function. If it’s not provided by the user, we’ll just have it replace the old state with the new when an update comes. The merge function can have any type in the argument for the final dispatch, but we assume T unless specified.

Now we should add the state that maintains the optimistic value.

export function useOptimistic<T>(
  initial: T,
  merge?: StateUpdater<T, A>,
): readonly [T, React.Dispatch<React.SetStateAction<A>>] {
  const [state, setState] = useState(initial)
  return [state, setState]
}

We… have a very simple hook working right now. It doesn’t really do anything special at all. Very bland.

Ok so, what’s the most basic requirement? Well, for starters, it should let us merge the state if we provide a merge function, which we haven’t use yet from our arguments. Then, we should take into account that updating the parent state (held in initial) should also update the state itself, so that it will be up-to-date with the latest server value.

Creating the update functions

We’ll start by adding the useEffect hook, that will update from the parent initial prop when it updates.

useEffect(() => {
  setState(initial)
}, [initial])

Now, when the initial value is updated, the internal state will be kept up-to-date with the change.

We can also implement the merge functionality. For that, we need to provide another implementation of the setState function, which uses it internally and adds some more logic.

const update = useCallback(
  (value: SetStateAction<A>) => {
    const newValue = typeof value === 'function' ? (value as (value: T) => T)(state) : value
    const merged = merge ? merge(state, newValue as A) : (newValue as T)
    setState(merged)
  },
  [state],
)

return [state, update] // instead of [state, setState]

Okay, it’s a little more to break down here.

First, we have an identical callback to setState now, except it also runs the merge function before sending the final value.

In useState, in addition to passing a new state value, the setState function can be used with a callback function, which takes the current state as a parameter, and should return the new state. The variable newValue is populated with the result of that process.

Then we just stack the merge function on top of that, if it exists.

The usage for this would still be:

const [count, setCount] = useOptimistic(0)

// option 1
setCount(1)

// option 2
setCount((c) => c + 1)

All the changes happen internally. Nice.

That’s pretty much it as it stands right now. React’s implementation is pretty similar to this, with maybe a slight variation on how the change is calculated with the parent’s update, to determine when the value should be brought up to date.

A bit anti-climactic? Well, let’s add a few bells and whistles. And we can now, because we don’t need React to build it for us. We’re strong!

What else can we add? Rolling back on failure

The first idea for improvement I have is to be able to catch an error, and roll back the value if it fails. Let’s begin implementing this.

First, we need to have a value to go back to. So every time the developer updates the optimistic state to a new value, and an operation related to it fails, we want to provide a way to revert to the original value.

For a recap, here’s the entire hook as it now is, before we make changes:

type StateUpdater<T, A = T> = (currentState: T, newState: A) => T

export function useOptimistic<T>(
  initial: T,
  merge?: StateUpdater<T, A>,
): readonly [T, React.Dispatch<React.SetStateAction<A>>] {
  const [state, setState] = useState(initial)

  useEffect(() => {
    setState(initial)
  }, [initial])

  const update = useCallback(
    (value: SetStateAction<A>) => {
      const newValue = typeof value === 'function' ? (value as (value: T) => T)(state) : value
      const merged = merge ? merge(state, newValue as A) : (newValue as T)
      setState(merged)
    },
    [state],
  )

  return [state, update]
}

Now, let’s say the developer comes across a case like this:

function MyApp({ comments }) {
  const [oComments, addComment] = useOptimistic(
    comments || [],
    (existingComments, newComment) => [...existingComments, newComment],
  )

  const sendComment = async (comment: Comment) => {
    addComment(comment)
    await sendCommentToApi(comment) // could fail and throw
  }

  return ...
}

In this case, we are sending some API request to send the comment content to the backend and make sure it’s already there.

But if it fails, we will have to revert it to its original content manually by keeping a separate state for the original value, and set it back if it failed.

Instead, we can make our new hook handle this.

Let’s add another state that maintains the previous value:

const [rollbackState, setRollbackState] = useState<T | null>(null)

And now that we have that, we should set this rollback value once we set a new value through the regular setter:

const update = useCallback(
  (value: SetStateAction<A>) => {
    const newValue = typeof value === 'function' ? (value as (value: T) => T)(state) : value
    const merged = merge ? merge(state, newValue as A) : (newValue as T)
    setRollbackState(state)
    setState(merged)
  },
  [state],
)

We take this value, save it, and… Do nothing with it. We should probably write that logic somewhere. How’s this:

const rollback = useCallback(() => {
  if (rollbackState == null) {
    // you can either return here, or throw an error, and be more strict about false rollbacks:
    // throw new Error('There is no state to roll back to')
    return
  }
  setState(rollbackState!)
  setRollbackState(null)
}, [rollbackState])

And we can also add a done callback, which ditches the retained value, just for clean up. Something like this:

const done = useCallback(() => {
  setRollbackState(null)
}, [])

Let’s also return these in our hook:

export function useOptimistic<T, A = T>(
  initial: T,
  merge?: StateUpdater<T, A>,
): readonly [
  T,
  React.Dispatch<React.SetStateAction<A>>,
  {
    readonly rollback: () => void
    readonly done: () => void
  },
] {
  // ...
  return [state, update, { rollback, done }]
}

Nice, right? We just revert. Now to actually use it, we can update the component:

const [oComments, addComment, { done, rollback }] = useOptimistic(
  //                          ^^^^^^^^^^^^^^^^^^
  comments || [],
  (existingComments, newComment) => [...existingComments, newComment],
)

const sendComment = async (comment: Comment) => {
  try {
    addComment(comment)
    await sendCommentToApi(comment) // could fail and throw
  } catch (e) {
    rollback()
  }
}

That’s not too bad. If we fail, we use rollback and get back our previous value, and we never had to save it separately, which is good for scaling this to multiple places.

But… if you give it a go, you will notice a problem.

Let’s do a small test:

const fail = async (comment: Comment) => {
  try {
    addComment(comment)
    await new Promise((res) => setTimeout(res, 1000))
    throw new Error('Oops!')
    done()
  } catch {
    rollback()
  }
}

Make this run somewhere and you will realize the value reverted back to null instead of the old value, or simply doesn’t roll back. What gives?

Here is the problem: when we are using rollback, it uses the state variable, rollbackState, to set the value back on the main state. But here is the thing - this variable is a dependency in the callback hook on rollback. Take a closer look:

// before:
// const [rollbackState, setRollbackState] = useState<T | null>(null)
//
// after:
const rollbackState = useRef<T | null>(null)

const rollback = useCallback(() => {
  if (rollbackState == null) {
    return
  }
  setState(rollbackState!)
  setRollbackState(null)
}, [rollbackState])
//  ^^^^^^^^^^^^^

So why is that bad? Well, in our fail test above, we used addComment which is our update function with a fancy name. When the API call fails we call rollback. But rollback was created bound to the rollbackState that was there at the time of its creation. We wait for a promise to return, and then use the done or rollback functions, that were created before the rollbackState was even set away from null and into a value in the first place. So basically we are re-using the value from before we put anything in it, which means it’s null, which is the value we roll back to when using rollback.

How do we fix this?

Fixing the stale value

There are probably many possible solutions, but here is my very naive one.

How do I make sure I get the correct reference to the updated value?

In reality, I don’t need rollbackState to be an actual state. Why can I say that with confidence? Well, I don’t need to call a re-render when it updates, as it is only there to reference a value in another function when called later, and it’s not displayed in the UI - meaning it has no need to call for a render when I update it. Also it happens on literally every update call as well, so I am trying to uselessly mark the frame dirty for re-render twice when one time is enough.

With that in mind, I can now safely switch the rollbackState from a state hook to a useRef hook.

Why would that work? Because a useRef reference is always a constant object. But the current value inside that object can be changed. In effect, when we reference a ref value, we can say we are sure it is the most recent, as the ref never needs to change and no mismatches can occur.

Here’s how it looks after the fix:

const update = useCallback(
  (value: SetStateAction<A>) => {
    const newValue = typeof value === 'function' ? (value as (value: T) => T)(state) : value
    const merged = merge ? merge(state, newValue as A) : (newValue as T)
    rollbackState.current = state
    setState(merged)
  },
  [state],
)

const rollback = useCallback(() => {
  if (rollbackState.current == null) {
    return
  }
  setState(rollbackState.current!)
  rollbackState.current = null
}, [])

const done = useCallback(() => {
  rollbackState.current = null
}, [])

This way, no matter when I call the rollback or done functions, they will always try to access the same parent object. With that object, they can modify the inner property and we don’t need to worry about the data being stale.

Simplifying usage

You might notice that now you can have a pretty redundant pattern of try {...} catch() {...} where you would rollback and done at the end of each section. Well that’s no fun, we want less code.

Why not add this to our hook?

const run = useCallback(
  async (cb: () => Promise<unknown> | unknown) => {
    try {
      await cb()
      done()
    } catch (_) {
      rollback()
    }
  },
  [done, rollback],
)

Simple yet elegant. If we export this from the hook:

return [state, update as React.Dispatch<React.SetStateAction<A>>, { done, rollback, run }]

Now it will now let us do:

const sendComment = async (comment: Comment) => {
  run(async () => {
    addComment(comment)
    await sendCommentToApi(comment)
  })
}

Much simpler, right?

Saving the promise state

One last thing I want us to add is the ability to know when our state is loading. This way, even without using useMutation from @tanstack/react-query, we can get a loading state of any promise, while working directly with the useOptimistic hook instead.

For that we can simply add a state or a ref (decide which you want depending on your wanted DX), which sets a pending state to true and false when we update the state, and when we call done/rollback, accordingly.

Here is the full entire hook with all the things we made here:

type StateUpdater<T, A = T> = (currentState: T, newState: A) => T

export function useOptimistic<T, A = T>(
  initial: T,
  merge?: StateUpdater<T, A>,
): readonly [
  T,
  React.Dispatch<React.SetStateAction<A>>,
  {
    readonly done: () => void
    readonly rollback: () => void
    readonly pending: boolean
    readonly run: (promise: () => Promise<unknown> | unknown) => void
  },
] {
  const [state, setState] = useState(initial)
  const rollbackState = useRef<T | null>(null)
  const [pending, setPending] = useState(false)

  useEffect(() => {
    setPending(false)
    setState(initial)
  }, [initial])

  const update = useCallback(
    (value: SetStateAction<A>) => {
      const newValue = typeof value === 'function' ? (value as (value: T) => T)(state) : value
      const merged = merge ? merge(state, newValue as A) : (newValue as T)
      rollbackState.current = state
      setState(merged)
      setPending(true)
    },
    [state],
  )

  const rollback = useCallback(() => {
    if (rollbackState.current == null) {
      // throw new Error('There is no state to roll back to')
      return
    }
    setState(rollbackState.current!)
    rollbackState.current = null
    setPending(false)
  }, [])

  const done = useCallback(() => {
    rollbackState.current = null
    setPending(false)
  }, [])

  const run = useCallback(
    async (cb: () => Promise<unknown> | unknown) => {
      try {
        await cb()
        done()
      } catch (_) {
        rollback()
      }
    },
    [done, rollback],
  )

  return [
    state,
    update as React.Dispatch<React.SetStateAction<A>>,
    { done, rollback, pending, run },
  ] as const
}

Where to go from here?

There are many ways that you can improve this hook, but I feel like it now reached a fairly usable state which we can all build from in our own special ways.

Until React decides (or decides against) promoting useOptimistic to non-canary channels, maybe we can all benefit from a relatively simple hook that does exactly this without having to wait.

I hope you found this article useful or interesting. Feel free to ask questions or suggest improvements in the comments!