Migrating Async/Await to Promise

.size-snapshot.json

@@ -77,9 +77,9 @@
     }
   },
   "middleware.js": {
-    "bundled": 5547,
-    "minified": 2945,
-    "gzipped": 1293,
+    "bundled": 6548,
+    "minified": 3299,
+    "gzipped": 1388,
     "treeshaked": {
       "rollup": {
         "code": 0,

src/middleware.ts

@@ -204,15 +204,52 @@ type PersistOptions<S> = {
   migrate?: (persistedState: any, version: number) => S | Promise<S>
 }
 
+interface Thenable<Value> {
+  then<V>(
+    onFulfilled: (value: Value) => V | Promise<V> | Thenable<V>
+  ): Thenable<V>
+  catch<V>(
+    onRejected: (reason: Error) => V | Promise<V> | Thenable<V>
+  ): Thenable<V>
+}
+
+const toThenable = <Result, Input>(
+  fn: (input: Input) => Result | Promise<Result> | Thenable<Result>
+) => (input: Input): Thenable<Result> => {
+  try {
+    const result = fn(input)
+    if (result instanceof Promise) {
+      return result as Thenable<Result>
+    }
+    return {
+      then(onFulfilled) {
+        return toThenable(onFulfilled)(result as Result)
+      },
+      catch(_onRejected) {
+        return this as Thenable<any>
+      },
+    }
+  } catch (e) {
+    return {
+      then(_onFulfilled) {
+        return this as Thenable<any>
+      },
+      catch(onRejected) {
+        return toThenable(onRejected)(e)
+      },
+    }
+  }
+}
+
 export const persist = <S extends State>(
   config: StateCreator<S>,
   options: PersistOptions<S>
 ) => (set: SetState<S>, get: GetState<S>, api: StoreApi<S>): S => {
   const {
     name,
     getStorage = () => localStorage,
-    serialize = JSON.stringify,
-    deserialize = JSON.parse,
+    serialize = JSON.stringify as (state: StorageValue<S>) => string,
+    deserialize = JSON.parse as (str: string) => StorageValue<S>,
     blacklist,
     whitelist,
     onRehydrateStorage,
@@ -241,7 +278,9 @@ export const persist = <S extends State>(
     )
   }
 
-  const setItem = async () => {
+  const thenableSerialize = toThenable(serialize)
+
+  const setItem = (): Thenable<void> => {
     const state = { ...get() }
 
     if (whitelist) {
@@ -253,7 +292,18 @@ export const persist = <S extends State>(
       blacklist.forEach((key) => delete state[key])
     }
 
-    return storage?.setItem(name, await serialize({ state, version }))
+    let errorInSync: Error | undefined
+    const thenable = thenableSerialize({ state, version })
+      .then((serializedValue) =>
+        (storage as StateStorage).setItem(name, serializedValue)
+      )
+      .catch((e) => {
+        errorInSync = e
+      })
+    if (errorInSync) {
+      throw errorInSync
+    }
+    return thenable
   }
 
   const savedSetState = api.setState
@@ -264,43 +314,61 @@ export const persist = <S extends State>(
   }
 
   // rehydrate initial state with existing stored state
-  ;(async () => {
-    const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined
-
-    try {
-      const storageValue = await storage.getItem(name)
 
+  // a workaround to solve the issue of not storing rehydrated state in sync storage
+  // the set(state) value would be later overridden with initial state by create()
+  // to avoid this, we merge the state from localStorage into the initial state.
+  let stateFromStorageInSync: S | undefined
+  const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined
+  // bind is used to avoid `TypeError: Illegal invocation` error
+  toThenable(storage.getItem.bind(storage))(name)
+    .then((storageValue) => {
       if (storageValue) {
-        const deserializedStorageValue = await deserialize(storageValue)
-
-        // if versions mismatch, run migration
+        return deserialize(storageValue)
+      }
+    })
+    .then((deserializedStorageValue) => {
+      if (deserializedStorageValue) {
         if (deserializedStorageValue.version !== version) {
-          const migratedState = await migrate?.(
-            deserializedStorageValue.state,
-            deserializedStorageValue.version
-          )
-          if (migratedState) {
-            set(migratedState)
-            await setItem()
+          if (migrate) {
+            return migrate(
+              deserializedStorageValue.state,
+              deserializedStorageValue.version
+            )
           }
+          console.error(
+            `State loaded from storage couldn't be migrated since no migrate function was provided`
+          )
         } else {
+          stateFromStorageInSync = deserializedStorageValue.state
           set(deserializedStorageValue.state)
         }
       }
-    } catch (e) {
+    })
+    .then((migratedState) => {
+      if (migratedState) {
+        stateFromStorageInSync = migratedState as S
+        set(migratedState as PartialState<S, keyof S>)
+        return setItem()
+      }
+    })
+    .then(() => {
+      postRehydrationCallback?.(get(), undefined)
+    })
+    .catch((e: Error) => {
       postRehydrationCallback?.(undefined, e)
-      return
-    }
-
-    postRehydrationCallback?.(get(), undefined)
-  })()
+    })
 
-  return config(
+  const configResult = config(
     (...args) => {
       set(...args)
       void setItem()
     },
     get,
     api
   )
+
+  return stateFromStorageInSync
+    ? { ...configResult, ...stateFromStorageInSync }
+    : configResult
 }

tests/persist.test.tsx

@@ -151,7 +151,7 @@ it('can migrate persisted state', async () => {
       name: 'test-storage',
       version: 13,
       getStorage: () => storage,
-      migrate: (state, version) => {
+      migrate: async (state, version) => {
         migrateCallCount++
         expect(state.count).toBe(42)
         expect(version).toBe(12)

tests/persistSync.test.tsx

@@ -0,0 +1,201 @@
+import create from '../src/index'
+import { persist } from '../src/middleware'
+
+const consoleError = console.error
+afterEach(() => {
+  console.error = consoleError
+})
+
+describe('persist middleware with sync configuration', () => {
+  it('can rehydrate state', () => {
+    let postRehydrationCallbackCallCount = 0
+
+    const storage = {
+      getItem: (name: string) =>
+        JSON.stringify({
+          state: { count: 42, name },
+          version: 0,
+        }),
+      setItem: () => {},
+    }
+
+    const useStore = create(
+      persist(
+        () => ({
+          count: 0,
+          name: 'empty',
+        }),
+        {
+          name: 'test-storage',
+          getStorage: () => storage,
+          onRehydrateStorage: () => (state, error) => {
+            postRehydrationCallbackCallCount++
+            expect(error).toBeUndefined()
+            expect(state?.count).toBe(42)
+            expect(state?.name).toBe('test-storage')
+          },
+        }
+      )
+    )
+
+    expect(useStore.getState()).toEqual({ count: 42, name: 'test-storage' })
+    expect(postRehydrationCallbackCallCount).toBe(1)
+  })
+
+  it('can throw rehydrate error', () => {
+    let postRehydrationCallbackCallCount = 0
+
+    const storage = {
+      getItem: () => {
+        throw new Error('getItem error')
+      },
+      setItem: () => {},
+    }
+
+    create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+        onRehydrateStorage: () => (_, e) => {
+          postRehydrationCallbackCallCount++
+          expect(e?.message).toBe('getItem error')
+        },
+      })
+    )
+
+    expect(postRehydrationCallbackCallCount).toBe(1)
+  })
+
+  it('can persist state', () => {
+    let setItemCallCount = 0
+
+    const storage = {
+      getItem: () => null,
+      setItem: (name: string, value: string) => {
+        setItemCallCount++
+        expect(name).toBe('test-storage')
+        expect(value).toBe(
+          JSON.stringify({
+            state: { count: 42 },
+            version: 0,
+          })
+        )
+      },
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        getStorage: () => storage,
+        onRehydrateStorage: () => (_, error) => {
+          expect(error).toBeUndefined()
+        },
+      })
+    )
+
+    expect(useStore.getState()).toEqual({ count: 0 })
+    useStore.setState({ count: 42 })
+    expect(useStore.getState()).toEqual({ count: 42 })
+    expect(setItemCallCount).toBe(1)
+  })
+
+  it('can migrate persisted state', () => {
+    let migrateCallCount = 0
+    let setItemCallCount = 0
+
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          state: { count: 42 },
+          version: 12,
+        }),
+      setItem: (_: string, value: string) => {
+        setItemCallCount++
+        expect(value).toBe(
+          JSON.stringify({
+            state: { count: 99 },
+            version: 13,
+          })
+        )
+      },
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        version: 13,
+        getStorage: () => storage,
+        onRehydrateStorage: () => (_, error) => {
+          expect(error).toBeUndefined()
+        },
+        migrate: (state, version) => {
+          migrateCallCount++
+          expect(state.count).toBe(42)
+          expect(version).toBe(12)
+          return { count: 99 }
+        },
+      })
+    )
+
+    expect(useStore.getState()).toEqual({ count: 99 })
+    expect(migrateCallCount).toBe(1)
+    expect(setItemCallCount).toBe(1)
+  })
+
+  it.only('can correclty handle a missing migrate function', () => {
+    console.error = jest.fn()
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          state: { count: 42 },
+          version: 12,
+        }),
+      setItem: (_: string, value: string) => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        version: 13,
+        getStorage: () => storage,
+        onRehydrateStorage: () => (_, error) => {
+          expect(error).toBeUndefined()
+        },
+      })
+    )
+
+    expect(useStore.getState()).toEqual({ count: 0 })
+    expect(console.error).toHaveBeenCalled()
+  })
+
+  it('can throw migrate error', () => {
+    let postRehydrationCallbackCallCount = 0
+
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          state: {},
+          version: 12,
+        }),
+      setItem: () => {},
+    }
+
+    const useStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        version: 13,
+        getStorage: () => storage,
+        migrate: () => {
+          throw new Error('migrate error')
+        },
+        onRehydrateStorage: () => (_, e) => {
+          postRehydrationCallbackCallCount++
+          expect(e?.message).toBe('migrate error')
+        },
+      })
+    )
+
+    expect(useStore.getState()).toEqual({ count: 0 })
+    expect(postRehydrationCallbackCallCount).toBe(1)
+  })
+})