Making Zustand Persist Play Nice with Async Storage & React Suspense, Part 2/2

Welcome to part 2 of my two part tutorial on Zustand persist with async storage and React Suspense. If you haven't seen part 1, you'll need the code from there, so check it out here. We left off having completed our store with the partial persist to IndexedDB and a merge function to ensure persisted data smoothly integrates with new store creation. Part 2 is about integrating Suspense and making our code even more robust by using onRehydratedStorage from Zustand to return data to the app only when the hydration has occurred. First of all, if you aren't using React Suspense I highly recommend it. It's actually incredibly easy to use, and by passing your fallback (I prefer skeletons unique to each component) React will automatically show it while the children passed to Suspense are loading. It also encourages use of lazy imports, since that is one of several ways you can trigger Suspense. Sometimes you might have a component that relies solely on data from the store and isn't using lazy, or lazy is working for the page but you need additional loading indicators for the data. So how would you trigger Suspense in that component? I'm going to show you by further customizing the code from part 1. This is going to be a mimic of the same effect you can get when you use, for example, useSuspenseQuery from Tanstack React Query. import { useSyncExternalStore} from 'react'; // NEW export const useStore = create()( persist( (set, get) => ({ // state bookmarks: {}, previousSearches: [], continueWatching: [], // suspense-related state THIS IS NEW isLoaded: false, isLoading: false, loadError: null, listeners: new Set void>(), // NEW method to subscribe to store changes (for useSyncExternalStore) subscribe: (listener) => { const { listeners } = get(); listeners.add(listener); return () => listeners.delete(listener); }, // NEW method to initialize the store and handle Suspense initializeStore: async () => { // Return immediately if already loaded if (get().isLoaded) { return { bookmarks: get().bookmarks, previousSearches: get().previousSearches, continueWatching: get().continueWatching, }; } // If already loading, wait for completion if (get().isLoading) { return new Promise((resolve, reject) => { const unsubscribe = get().subscribe(() => { if (get().isLoaded) { unsubscribe(); resolve({ bookmarks: get().bookmarks, previousSearches: get().previousSearches, continueWatching: get().continueWatching, }); } else if (get().loadError) { unsubscribe(); reject(get().loadError); } }); }); } // Start loading process set({ isLoading: true }); get().listeners.forEach((listener) => listener()); // Return a promise that resolves when loaded return new Promise((resolve, reject) => { const unsubscribe = get().subscribe(() => { if (get().isLoaded) { unsubscribe(); resolve({ bookmarks: get().bookmarks, previousSearches: get().previousSearches, continueWatching: get().continueWatching, }); } else if (get().loadError) { unsubscribe(); reject(get().loadError); } }); }); }, // other state updating methods go here ... }), { name: process.env.NODE_ENV === 'production' ? 'idb-storage' : 'idb-storage-dev', storage: idbStorage, version: 0, merge: (persistedState, currentState) => { .... }, //NEW - how we know when hydration is done onRehydrateStorage: () => (_state, error) => { if (error) { console.log('An error occurred during hydration', error); useStore.setState({ isLoaded: true, isLoading: false, loadError: error as Error, }); } else { useStore.setState({ isLoaded: true, isLoading: false, }); } useStore.getState().listeners.forEach((listener) => listener()); }, // other .... } ) ); // hook to trigger the initializeStore and wait for data while triggering Suspense: export function useSuspenseStore(selector: (_state: BingeBoxStore) => T): T { const store = useStore(); // if not loaded and not loading, start the loading process if (!store.isLoaded && !store.isLoading) { // This will be caught by React Suspense throw store.initializeStore(); } // If currently loading, throw a promise to trigger suspense if (st

Apr 22, 2025 - 07:49

Making Zustand Persist Play Nice with Async Storage & React Suspense, Part 2/2

Welcome to part 2 of my two part tutorial on Zustand persist with async storage and React Suspense. If you haven't seen part 1, you'll need the code from there, so check it out here.

We left off having completed our store with the partial persist to IndexedDB and a merge function to ensure persisted data smoothly integrates with new store creation. Part 2 is about integrating Suspense and making our code even more robust by using onRehydratedStorage from Zustand to return data to the app only when the hydration has occurred.

First of all, if you aren't using React Suspense I highly recommend it. It's actually incredibly easy to use, and by passing your fallback (I prefer skeletons unique to each component) React will automatically show it while the children passed to Suspense are loading. It also encourages use of lazy imports, since that is one of several ways you can trigger Suspense. Sometimes you might have a component that relies solely on data from the store and isn't using lazy, or lazy is working for the page but you need additional loading indicators for the data. So how would you trigger Suspense in that component? I'm going to show you by further customizing the code from part 1. This is going to be a mimic of the same effect you can get when you use, for example, useSuspenseQuery from Tanstack React Query.

import { useSyncExternalStore} from 'react'; // NEW

export const useStore = create()(
  persist(
    (set, get) => ({
      // state
      bookmarks: {},
      previousSearches: [],
      continueWatching: [],

      // suspense-related state THIS IS NEW
      isLoaded: false,
      isLoading: false,
      loadError: null,
      listeners: new Set<() => void>(),

      // NEW method to subscribe to store changes (for useSyncExternalStore)
      subscribe: (listener) => {
        const { listeners } = get();
        listeners.add(listener);
        return () => listeners.delete(listener);
      },

      // NEW method to initialize the store and handle Suspense
      initializeStore: async () => {
        // Return immediately if already loaded
        if (get().isLoaded) {
          return {
            bookmarks: get().bookmarks,
            previousSearches: get().previousSearches,
            continueWatching: get().continueWatching,
          };
        }

        // If already loading, wait for completion
        if (get().isLoading) {
          return new Promise((resolve, reject) => {
            const unsubscribe = get().subscribe(() => {
              if (get().isLoaded) {
                unsubscribe();
                resolve({
                  bookmarks: get().bookmarks,
                  previousSearches: get().previousSearches,
                  continueWatching: get().continueWatching,
                });
              } else if (get().loadError) {
                unsubscribe();
                reject(get().loadError);
              }
            });
          });
        }

        // Start loading process
        set({ isLoading: true });
        get().listeners.forEach((listener) => listener());

        // Return a promise that resolves when loaded
        return new Promise((resolve, reject) => {
          const unsubscribe = get().subscribe(() => {
            if (get().isLoaded) {
              unsubscribe();
              resolve({
                bookmarks: get().bookmarks,
                previousSearches: get().previousSearches,
                continueWatching: get().continueWatching,
              });
            } else if (get().loadError) {
              unsubscribe();
              reject(get().loadError);
            }
          });
        });
      },

           // other state updating methods go here ...
    }),
{
      name:
        process.env.NODE_ENV === 'production'
          ? 'idb-storage'
          : 'idb-storage-dev',
      storage: idbStorage,
      version: 0,
      merge: (persistedState, currentState) => {
                       ....
      },
//NEW - how we know when hydration is done
      onRehydrateStorage: () => (_state, error) => {
        if (error) {
          console.log('An error occurred during hydration', error);
          useStore.setState({
            isLoaded: true,
            isLoading: false,
            loadError: error as Error,
          });
        } else {
          useStore.setState({
            isLoaded: true,
            isLoading: false,
          });
        }

        useStore.getState().listeners.forEach((listener) => listener());
      },

// other ....

 }
  )
);

// hook to trigger the initializeStore and wait for data while triggering Suspense:
export function useSuspenseStore(selector: (_state: BingeBoxStore) => T): T {
  const store = useStore();

  // if not loaded and not loading, start the loading process
  if (!store.isLoaded && !store.isLoading) {
    // This will be caught by React Suspense
    throw store.initializeStore();
  }

  // If currently loading, throw a promise to trigger suspense
  if (store.isLoading) {
    throw new Promise((resolve) => {
      const unsubscribe = store.subscribe(() => {
        if (store.isLoaded || store.loadError) {
          unsubscribe();
          resolve(null);
        }
      });
    });
  }

  // If there was an error, throw it
  if (store.loadError) {
    throw store.loadError;
  }

  // Store is loaded, use it
  return useSyncExternalStore(store.subscribe, () => selector(store));
}

uSES is the key here, allowing us to integrate Zustand into React so that React actually knows about it. Bring that in first, then you should recognize state from part 1, but additional state is now keeping track of loading status, error, and listeners. Two new methods include subscribe code for uSES and initializeStore which will return a Promise that can be thrown by the useSuspenseStore hook to trigger Suspense. onRehydrateStorage is from the Zustand library and lets us properly set loading state based on hydration being complete. If hydration is complete, isLoading is false and isLoaded is true.

Anytime you don't need to trigger Suspense and you just want to use your store selector, just use the useStore hook that was created with this line of code:
export const useStore = create()(
That will likely be the usual way you access your store. This code is for when your 'get' functionality to your store needs to integrate with Suspense to show fallbacks.

Bonus: From the Zustand docs, one of the costs of using async storage is that async hydration might not be ready when React renders. You can use the isLoaded flag if you need to hold back rendering until some data is hydrated. You can also just use onRehydrateStorage by itself to know when the store is hydrated. The only reason for the initiateStorage and uSES was to integrate Suspense.

If you want to know about uSES and other ways it can be useful check out this blog.

Thanks for reading and I hope this helped in your understanding of some more advanced concepts of Zustand, Suspense, and async storage. If you have any comments or questions please let me know. If you notice that I am overlooking something I would love to hear from you! This is my own custom solution to this problem so any critique would be welcome. I have purposefully omitted showing the React code for usage of these hooks but that is available in React docs or Zustand docs if you are unsure how to use hooks.