Zustand and React Context

Zustand is a great lib for global client-state management. It's simple, fast, and has a small bundle size. There is however one thing I don't necessarily like about it:

The stores are global.

Okay? But isn't that the point of global state management? To make that state available in your app, everywhere?

Sometimes, I think that's true. However, as I've looked at how I've been using zustand over the last couple of years, I've realized that more often than not, I've needed some state to be available globally to one component subtree rather than the whole application. With zustand, it's totally fine - encouraged even - to create multiple, small stores on a per-feature basis. So why would I need my Dashboard Filters store to be available globally if I only need it on the Dashboard route? Sure, I can do that when it doesn't hurt, but I've found that having global stores do have a couple of drawbacks:

Initializing from Props

Global stores are created outside of the React Component lifecycle, so we can't initialize the store with a value we get as a prop. With a global store, we need to create it first with a known default state, then sync props-to-store with useEffect:

const useBearStore = create((set) => ({
  // ⬇️ initialize with default value
  bears: 0,
  actions: {
    increasePopulation: (by) =>
      set((state) => ({ bears: state.bears + by })),
    removeAllBears: () => set({ bears: 0 }),
  },
}))

const App = ({ initialBears }) => {
  //😕 write initialBears to our store
  React.useEffect(() => {
    useBearStore.set((prev) => ({ ...prev, bears: initialBears }))
  }, [initialBears])

  return (
    <main>
      <RestOfTheApp />
    </main>
  )
}

Apart from not wanting to write useEffect, this isn't ideal for two reasons:

1. We first render <RestOfTheApp /> with bears: 0 before the effect kicks in, then render it once more with the correct initialBears.
2. We don't really initialize our store with initialBears - we sync it. So if the initialBears change, we will see the update reflected in our store as well.

Testing Challenges

I find the testing docs for zustand pretty confusing and complex. It's all about mocking zustand and resetting the store between tests and so on. I think it all stems from the fact that the store is global. If it were scoped to a component subtree, we could render those components and the store would be isolated to it, not needing any of those "workarounds".

Reusability Limitations

Not all stores are singletons that we can use once in our App or once in a specific route. Sometimes, we want zustand stores for reusable components as well. One example from the past I can think of is a complex, multi-selection group component from our design-system. It was using local state passed down with React Context to manage the internal state of the selections. It became sluggish whenever an item was selected as soon as there were fifty or more items.

The Solution: React Context

Interestingly, there is a single way to fix all of these problems: React Context.

It's funny and ironic that React Context is the solution here, because using Context as a state management tool is what caused the aforementioned issue in the first place. But that's not what I'm proposing. The idea is to merely share the store instance via React Context - not the store values themselves.

Here's how that can look like with v5 syntax:

import { createStore, useStore } from 'zustand'

const BearStoreContext = React.createContext(null)

const BearStoreProvider = ({ children, initialBears }) => {
  const [store] = React.useState(() =>
    createStore((set) => ({
      bears: initialBears,
      actions: {
        increasePopulation: (by) =>
          set((state) => ({ bears: state.bears + by })),
        removeAllBears: () => set({ bears: 0 }),
      },
    }))
  )

  return (
    <BearStoreContext.Provider value={store}>
      {children}
    </BearStoreContext.Provider>
  )
}

The main difference here is that we aren't using create like before, which would give us a ready-to-use hook. Instead, we are relying on the vanilla zustand function createStore, which will just create a store for us.

Creating a custom hook to consume the store:

const useBearStore = (selector) => {
  const store = React.useContext(BearStoreContext)
  if (!store) {
    throw new Error('Missing BearStoreProvider')
  }
  return useStore(store, selector)
}

Benefits of this Approach

This solution solves three key problems:

1. Initialization from Props: We can now initialize our store with props by creating it inside the React Component tree.
2. Easy Testing: Store creation becomes isolated to specific tests, eliminating the need for complex mocking.
3. Reusability: Components can now have their own encapsulated zustand stores, enabling true component-level state management.

So even though the zustand docs pride themselves about not needing a Context Provider to access a store, I think knowing how to combine store creation with React Context can come in quite handy in situations where encapsulation and reusability are required.