Integrating Zustand with next.js app router

Integrating Zustand with next.js app router
Zustand integration with Next.js

Zustand is a popular state management library in the React ecosystem. It's often preferred over Redux because it's so simple and easy to learn. However, integrating Zustand into a Next.js app router project requires a specific approach.

WHY?

To understand why, we first need to touch on the Singleton pattern. This is a design pattern that ensures a class has only one instance and provides a single, global way to access it. The basic idea is to check if an instance of the class already exists. If it does, you simply reuse that same instance. If it doesn't, you create a new one, save it for later, and then return it.

const Singleton = (function() {
    let instance;

    function createInstance() {
        const object = new Object("I am the single instance");
        return object;
    }

    return {
        getInstance: function() {
            if (!instance) {
                instance = createInstance();
            }
            return instance;
        }
    };
})();

// Get the instance for the first time
const instance1 = Singleton.getInstance();

// Get the instance again
const instance2 = Singleton.getInstance();

// Check if both instances are the same
console.log(instance1 === instance2); // true
console.log(instance1); // Object { "I am the single instance" }

By default, JavaScript modules function similarly to a singleton, even though they aren't a strict implementation of the classic design pattern. The first time a module is imported, its code runs and its exports are cached. Any future imports of that same module will simply return a reference to that cached object, instead of running the code again and creating a new instance. This means that any values or objects exported from a module are shared across all parts of your application that import it.

So, when Zustand is used in a normal React SPA environment, it follows the same pattern. The store is created once, and that same object is used everywhere. This is the reason why we need to create a separate Zustand store for each request while using Zustand in Next.js.

Since a Next.js server can handle multiple requests simultaneously, a global store (which is the default behavior of Zustand) would be shared across all of them. This would lead to a serious security and data integrity issue. In a client-side only React app, this is fine because each user has their own instance of the application running in their browser.

HOW ?

To implement a per request store for Next.js, we can use react context and useRef hook.

'use client';

import { createContext, useContext, useRef } from 'react';
import { createStore, useStore } from 'zustand';

// 1. Define the store's state and actions
interface CounterState {
  count: number;
  increment: () => void;
  decrement: () => void;
}

// 2. Define the store's creation function.
// This function will be used to create a new store instance for each request.
const createCounterStore = () =>
  createStore<CounterState>((set) => ({
    count: 0,
    increment: () => set((state) => ({ count: state.count + 1 })),
    decrement: () => set((state) => ({ count: state.count - 1 })),
  }));

// 3. Create a React Context to hold the store instance
const CounterStoreContext = createContext<ReturnType<typeof createCounterStore> | null>(null);

// 4. The StoreProvider component. This is where the magic happens.
// It creates a new store instance per request and provides it to its children.
export function CounterStoreProvider({ children }: { children: React.ReactNode }) {
  // Use useRef to create the store instance only once per component instance.
  // This is essential for preventing the store from being recreated on every re-render.
  const storeRef = useRef<ReturnType<typeof createCounterStore>>();
  if (!storeRef.current) {
    storeRef.current = createCounterStore();
  }

  return (
    <CounterStoreContext.Provider value={storeRef.current}>
      {children}
    </CounterStoreContext.Provider>
  );
}

// 5. A custom hook to access the store from any component within the provider's scope.
const useCounterStore = <T,>(selector: (state: CounterState) => T): T => {
  const store = useContext(CounterStoreContext);
  if (!store) {
    throw new Error('useCounterStore must be used within a CounterStoreProvider');
  }
  return useStore(store, selector);
};

And then, we use that store in a client component that is within the zustand provider

'use client';

const Counter = () => {
  const count = useCounterStore((state) => state.count);
  const increment = useCounterStore((state) => state.increment);
  const decrement = useCounterStore((state) => state.decrement);

  return (
    <div className="flex flex-col items-center justify-center p-8 bg-gray-100 dark:bg-gray-800 rounded-lg shadow-lg space-y-4">
      <h2 className="text-2xl font-bold text-gray-800 dark:text-gray-100">Zustand Counter</h2>
      <p className="text-6xl font-extrabold text-indigo-600 dark:text-indigo-400">{count}</p>
      <div className="flex space-x-4">
        <button
          onClick={increment}
          className="px-6 py-3 text-lg font-semibold text-white bg-indigo-500 rounded-lg shadow-md hover:bg-indigo-600 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-opacity-50 transition duration-300 ease-in-out"
        >
          Increment
        </button>
        <button
          onClick={decrement}
          className="px-6 py-3 text-lg font-semibold text-white bg-red-500 rounded-lg shadow-md hover:bg-red-600 focus:outline-none focus:ring-2 focus:ring-red-500 focus:ring-opacity-50 transition duration-300 ease-in-out"
        >
          Decrement
        </button>
      </div>
    </div>
  );
};

This approach ensures that a new, isolated store is created for each server request, preventing state from being shared between different users. The useRef hook is crucial here because it ensures the store instance is created only once during the component's lifetime, which is important for preventing unnecessary re-renders and maintaining a consistent store.

Notice the use client directive at top. We need to use that because we're using context and hooks that don't work in a Next.js server-side component. So we need to make it a client component.

You might ask if we use Zustand only inside client components (use client), can we then use a single global Zustand store? The answer is no, we can't. Because even for client components, the initial rendering happens in Next.js server. So the data integrity and security risks are still there.

Another question you might ask is whether using context to provide the Zustand store will cause a re-render of the whole subsequent tree anytime the store is updated?

The answer is yet another No. When you create a Zustand store, you're not actually getting a simple object with state properties. You're getting a hook (or a function) that, behind the scenes, manages an internal store object. This internal store object is the stable reference we've been talking about. Here is a simplified example of Zustand store object:

// This is what you define with `create()`
const createStore = (set, get) => ({
  count: 0,
  increment: () => set(state => ({ count: state.count + 1 })),
  // ... other state properties and actions
});

// This is the "stable" store object that Zustand creates and returns a hook to.
// The `current` property is what gets replaced.
const zustandStoreObject = {
  // A property to hold the current state object.
  // This is what gets replaced every time you call `set`.
  current: {
    count: 0,
    increment: () => {
      // Logic that updates the `current` property
      // with a new state object.
    },
  },

  // A list of all components that are subscribed to the store.
  // This is how Zustand knows which components to re-render.
  subscribers: new Set(),
  
  // Methods to interact with the store.
  getState: () => zustandStoreObject.current,
  setState: (newState) => {
    // This is the core of the immutability.
    // It creates a *new* state object and replaces the old one.
    zustandStoreObject.current = { ...zustandStoreObject.current, ...newState };
    // Then it notifies all the subscribers to re-render.
    zustandStoreObject.subscribers.forEach(subscriber => subscriber());
  },
  subscribe: (callback) => zustandStoreObject.subscribers.add(callback),
  unsubscribe: (callback) => zustandStoreObject.subscribers.delete(callback),
};

When you use the hook const useStore = create(...), you're getting a function that gives you access to this zustandStoreObject. When you call const count = useStore(state => state.count), the hook does two things:

  1. It subscribes your component to the zustandStoreObject.
  2. It returns zustandStoreObject.current.count.

In our Next.js setup, we pass this zustandStoreObject through a React Context, the reference to this object never changes. The current property inside it changes, but the useContext hook doesn't care about that. It only cares if the object it's watching is a new object. Zustand's internal subscription mechanism is what handles the updates and re-renders, making it efficient.