Zustand persist middleware unexpectedly saves state retrieved from persist layer

[double-thinker]

Bug Description

When using Zustand's persist middleware, I've noticed an unexpected behavior. The middleware is saving the state to storage even when that state was just retrieved from the persist layer even when no migration happens. This seems redundant and could be detrimental for some kind of storages.

Example:

// index.mjs
import { createStore } from "zustand/vanilla";
import { persist } from 'zustand/middleware'

const store = createStore(persist((set, get) => ({
    a: "initial"
}), {
    name: 'test',
    storage: {
        getItem: (_key) => {
            console.log("Getting");
            return { state: { a: "hydrated" } }
        },
        setItem: (_key, value) => {
            console.log("Setting", value)
        },
    },
    onRehydrateStorage: (state) => {
        console.log('hydration starts')

        // optional
        return (state, error) => {
            if (error) {
                console.log('an error happened during hydration', error)
            } else {
                console.log('hydration finished')
            }
        }
    }
}));

console.log("State:", store.getState());

Output:

hydration starts
Getting
Setting { state: { a: 'hydrated' }, version: 0 }
hydration finished
State: { a: 'hydrated' }

Expected behavior: The persist middleware should not write back to storage if no migration happens. A simple check for migration could be added to prevent triggering setItem() unnecessarily:

diff --git a/src/middleware/persist.ts b/src/middleware/persist.ts
index fc86500..87707f7 100644
--- a/src/middleware/persist.ts
+++ b/src/middleware/persist.ts
@@ -303,27 +303,33 @@ const oldImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
             deserializedStorageValue.version !== options.version
           ) {
             if (options.migrate) {
-              return options.migrate(
-                deserializedStorageValue.state,
-                deserializedStorageValue.version,
-              )
+              return [
+                true,
+                options.migrate(
+                  deserializedStorageValue.state,
+                  deserializedStorageValue.version,
+                ),
+              ]
             }
             console.error(
               `State loaded from storage couldn't be migrated since no migrate function was provided`,
             )
           } else {
-            return deserializedStorageValue.state
+            return [false, deserializedStorageValue.state]
           }
+        } else {
+          return [false, undefined]
         }
       })
-      .then((migratedState) => {
+      .then((migrationResult) => {
+        const [migrated, migratedState] = migrationResult as [boolean, S]
         stateFromStorage = options.merge(
           migratedState as S,
           get() ?? configResult,
         )
 
         set(stateFromStorage as S, true)
-        return setItem()
+        if (migrated) return setItem()
       })
       .then(() => {
         postRehydrationCallback?.(stateFromStorage, undefined)
@@ -459,27 +465,34 @@ const newImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
             deserializedStorageValue.version !== options.version
           ) {
             if (options.migrate) {
-              return options.migrate(
-                deserializedStorageValue.state,
-                deserializedStorageValue.version,
-              )
+              return [
+                true,
+                options.migrate(
+                  deserializedStorageValue.state,
+                  deserializedStorageValue.version,
+                ),
+              ]
             }
             console.error(
               `State loaded from storage couldn't be migrated since no migrate function was provided`,
             )
           } else {
-            return deserializedStorageValue.state
+            return [false, deserializedStorageValue.state]
           }
         }
+        return [false, undefined]
       })
-      .then((migratedState) => {
+      .then((migrationResult) => {
+        const [migrated, migratedState] = migrationResult as [boolean, S]
         stateFromStorage = options.merge(
           migratedState as S,
           get() ?? configResult,
         )
 
         set(stateFromStorage as S, true)
-        return setItem()
+        if (migrated) {
+          return setItem()
+        }
       })
       .then(() => {
         // TODO: In the asynchronous case, it's possible that the state has changed

Reproduction Link

https://stackblitz.com/edit/vitejs-vite-skujfr?file=index.mjs

[dai-shi]

Thanks for reporting! It sounds reasonable. I wonder how it causes an actual problem, but I understand it could.

The implementation looks good overall:

• Please only modify newImpl
• return [...] as const should avoid as [boolean, S]

Would you like to send a PR? A test would be necessary.

[double-thinker]

Perfect. I will submit a pull request in the coming days.

Thanks for all your work. Zustand is an amazing tool 😄

[double-thinker]

Edit: It was a typo from my side. as const works as you suggested.

@dai-shi Using returns [...] as const shows the following: Type 'readonly [true, T | Promise<T>] | readonly [false, T] | readonly [false, undefined] | undefined' must have a '[Symbol.iterator]()' method that returns an iterator.ts(2488). How do you prefer to handle this?

Btw, in my case this extra setItem all adds networks load because the storage is sync in a P2P fashion using yjs.

[dai-shi]

Thanks for the PR! #2678

the storage is sync in a P2P fashion using yjs.

Niceee.