onRehydrateStorage is called before store.getItem

[jeffreyyoung]

Summary
These are the logs for the below store

onRehydrateStorage {count: 0, _hasHydrated: false, increment: ƒ, setHasHydrated: ƒ}
storage.setItem food-storage {"state":{"_hasHydrated":true},"version":0}
storage.getItem food-storage
storage.getItem value {"state":{"_hasHydrated":true},"version":0}
storage.setItem food-storage {"state":{"count":0,"_hasHydrated":true},"version":0}

Because onRehydrateStorage is called before storage.getItem, the store is always reset with default data.

For the following store:

export const useCount = create()(
  persist(
    (set, get) => ({
      count: 0,
      _hasHydrated: false,
      increment: () => {
        set({ count: get().count + 1 });
      },
      setHasHydrated: (value: boolean) => {
        set({ _hasHydrated: value });
      },
    }),
    {
      name: "food-storage", // name of item in the storage (must be unique)
      storage: createJSONStorage(() => ({
        getItem: async (name) => {
          console.log("storage.getItem", name);
          const value = localStorage.getItem(name);
          console.log("storage.getItem value", value);
          return value;
        },
        removeItem: async (name) => localStorage.removeItem(name),
        setItem: async (name: string, value) => {
          console.log("storage.setItem", name, value);
          return localStorage.setItem(name, value);
        },
      })),
      // (optional) by default the 'localStorage' is used
      onRehydrateStorage: (data: any) => {
        console.log("onRehydrateStorage", data);
        data.setHasHydrated(true);
      },
    }
  )
);

Link to reproduction
https://codesandbox.io/p/sandbox/zustand-repro-pyys3x - see App.tsx and look at console logs

[dai-shi]

I wonder if anyone can help investigating it.

[dbritto-dev]

@jeffreyyoung would you mind creating the custom storage outside?

[jeffreyyoung]

I tried with custom storage outside but had the same result.

Here's the updated code:

import "./styles.css";
import { create } from "zustand";
import { persist, createJSONStorage } from "zustand/middleware";
interface MyState {
  count: number;
  increment: () => void;
  _hasHydrated: boolean;
  setHasHydrated: (value: boolean) => void;
}

var storage = createJSONStorage(() => ({
  getItem: async (name) => {
    console.log("storage.getItem", name);
    const value = localStorage.getItem(name);
    console.log("storage.getItem value", value);
    return value;
  },
  removeItem: async (name) => {
    console.log("remove", name);
    localStorage.removeItem(name);
  },
  setItem: async (name: string, value) => {
    console.log("storage.setItem", name, value);
    return localStorage.setItem(name, value);
  },
}));

export const useCount = create()(
  persist(
    (set, get) => ({
      count: 0,
      _hasHydrated: false,
      increment: () => {
        set({ count: get().count + 1 });
      },
      setHasHydrated: (value: boolean) => {
        set({ _hasHydrated: value });
      },
    }),
    {
      name: "food-storage", // name of item in the storage (must be unique)
      storage,
      onRehydrateStorage: (data: any) => {
        console.log("onRehydrateStorage", data);
        data.setHasHydrated(true);
      },
    }
  )
);
export default function App() {
  const store = useCount();
  return (
    <div className="App">
      <h2>{store.count}</h2>
      <button onClick={store.increment}>increment</button>
    </div>
  );
}

As a work around I use the "merge" param to set _hasHydrated

      merge(persisted, current) {

          return Object.assign({}, persisted, current, { _hasHydrated: true });
      },

[Dragollla]

it happens because of set() in onRehydrateStorage listener. Correct way would be to return a function from listener and update flag in it. like so:

      onRehydrateStorage: (data: any) => {
        console.log("onRehydrateStorage started", data);
        return () => data.setHasHydrated(true);
      },