Learning Subscription with useEffect from a Simplified Jotai Implementation

Hi,

It might come as a bit of a surprise, but Jotai does not use useSyncExternalStore as of now. You can learn more about this in my previous blog post:

Why useSyncExternalStore Is Not Used in Jotai

Today's post is about how to handle subscriptions with useEffect. Let's look at a simplified implementation of Jotai.

Reference: Jotai Core Internals

import { useState, useEffect } from 'react'

// atom function returns a config object which contains initial value
export const atom = (initialValue) => ({ init: initialValue })

// we need to keep track of the state of the atom.
// we are using weakmap to avoid memory leaks
const atomStateMap = new WeakMap()
const getAtomState = (atom) => {
  let atomState = atomStateMap.get(atom)
  if (!atomState) {
    atomState = { value: atom.init, listeners: new Set() }
    atomStateMap.set(atom, atomState)
  }
  return atomState
}

// useAtom hook returns a tuple of the current value
// and a function to update the atom's value
export const useAtom = (atom) => {
  const atomState = getAtomState(atom)
  const [value, setValue] = useState(atomState.value)
  useEffect(() => {
    const callback = () => setValue(atomState.value)

    // same atom can be used at multiple components, so we need to
    // keep listening for atom's state change till component is unmounted.
    atomState.listeners.add(callback)
    callback()
    return () => atomState.listeners.delete(callback)
  }, [atomState])

  const setAtom = (nextValue) => {
    atomState.value = nextValue

    // let all the subscribed components know that the atom's state has changed
    atomState.listeners.forEach((l) => l())
  }

  return [value, setAtom]
}

Everything should look straightforward, but notice the immediate invocation of callback inside the useEffect body. This line is not commented in the code, yet it requires an explanation.

It is a little tricky, but it is a well-known technique among state management library authors. The reason it exists is that the external store, which in this case is atomStateMap, can be modified before the effect fires but after the state is initialized. Since useEffect runs asynchronously, there is a chance that the external store may be mutated between the render and the effect execution. The callback() line ensures that the atom value is updated even if such a mutation happens. Now, I hope you understand why it is required.

Should You Use useEffect for Subscriptions?

This behavior only applies to subscription mechanisms based on useState and useEffect. The recommended approach for subscriptions is now the useSyncExternalStore hook, provided by React itself. It solves edge cases that even a useEffect-based solution cannot handle.

So, in most cases, you will not need useEffect for subscriptions. Nevertheless, I hope this explanation helps you understand the code. You may not need it, but at least you know why it is tricky.

Happy coding.