Repository: pmndrs/zustand
Pull request #403 (2021-05-26 17:32:23)

Migrating Async/Await to Promise
# WIP
`middleware.ts` is completely migrated from Async/Await to Promise.
and also sync apis will be called in a sync manner.
- [x] a general API created to handle both sync/async peacefully
- [x] async tests passed
- [x] sync tests added
- [x] sync test passed
- [x] types added

fixing #346

codesandbox[bot]: This pull request is automatically built and testable in [CodeSandbox](https://codesandbox.io).

  To see build info of the built libraries, click [here](https://ci.codesandbox.io/status/pmndrs/zustand/pr/403/builds/134053) or the icon next to each commit SHA.

Latest deployment of this branch, based on commit 658e9c99d331ce859b183edb11e3e5ca09ca6655:

|Sandbox| Source |
|--|--|
|[React](https://codesandbox.io/s/react-3ljme)| Configuration |
|[React Typescript](https://codesandbox.io/s/react-typescript-5idxd)| Configuration |

barbogast: @dai-shi https://gist.github.com/barbogast/5679eabc25b761da0f500f53ca3ea0c2 contains an idea for a fix of the remaining issue with the rehydration. Is this approach correct, or do you think it should be fixed in an entirely different way?
barbogast: Just updated both gists. The tests should be finished, and are passing against the updated `middleware.ts` gist. How do you want me to submit the code? A PR against https://github.com/TheEhsanSarshar/zustand/tree/async-await-fix? Or do you just want to grab the gists and add them to PR directly, @TheEhsanSarshar?
barbogast: I'm wondering... is it possible to correctly type `makeThenable()`? I had a quick try and didn't manage, especially since the returned object needs to be compatible with real Promises for TS (I guess?). 

That together with the subtle but important differences between sync mode and async mode within the same code block make me wonder if it's worth it to handle both cases with the same code. Maybe it would be better to duplicate the code and handle both cases separately. Not sure, though, just sharing the thought.
TheEhsanSarshar: > Just updated both gists. The tests should be finished, and are passing against the updated `middleware.ts` gist. How do you want me to submit the code? A PR against https://github.com/TheEhsanSarshar/zustand/tree/async-await-fix? Or do you just want to grab the gists and add them to PR directly, @TheEhsanSarshar?

I am going to make some changes to add support for sync error handling. then you can make a PR on the repo
TheEhsanSarshar: Hi dear @dai-shi.
in this gist, I just implemented a new `toThroable` that is able to handle errors in both syn/async. it also handles thrown exception inside `then` and `catch`
can you please have a look at that?
https://gist.github.com/TheEhsanSarshar/d9852a7cc102dae0f07186e1dc4fe700
TheEhsanSarshar: the async tests are failing. dear @barbogast would you mind please check that.
dai-shi: https://gist.github.com/TheEhsanSarshar/d9852a7cc102dae0f07186e1dc4fe700#gistcomment-3761417
The conversations are scattered.
We could just paste the code here.
barbogast: FWIW, I created a PR with my results agains @TheEhsanSarshar 's branch: https://github.com/TheEhsanSarshar/zustand/pull/1. The changes are separated by commit, so you can cherry-pick. It contains the tests for the sync storage and code adjustments to `middleware.ts` so that all the tests are passing.
barbogast: >in this gist, I just implemented a new toThenable that is able to handle errors in both syn/async. it also handles thrown exception inside then and catch
can you please have a look at that?
https://gist.github.com/TheEhsanSarshar/d9852a7cc102dae0f07186e1dc4fe700

Hmm... just realised, that it might be worth it to have separate tests `toThenable`, then. There are a lot of ways in which this function can be used, and without tests it might be tricky to guarantee that all of them are handled correctly. Especially thinking about all the different ways in which errors could occur.
dai-shi: Thanks guys for working on this.
I thought too that tests for `toThenable` would be nice.
But, we haven't agreed on the api yet.
Let me put my proposal here:
```js
const toThenable = (fn) => (input) => {
  try {
    const result = fn(input)
    if (result instanceof Promise) {
      return result
    }
    return {
      then(onFulfilled) { return toThenable(onFulfilled)(result) },
      catch(onRejected) { return this },
    }
  } catch (e) {
    return {
      then(onFulfilled) { return this },
      catch(onRejected) { return toThenable(onRejected)(e) },
    }
  }
}
```
TheEhsanSarshar: thanks @dai-shi I am going to playaround with your proposal 
TheEhsanSarshar: > > in this gist, I just implemented a new toThenable that is able to handle errors in both syn/async. it also handles thrown exception inside then and catch
> > can you please have a look at that?
> > https://gist.github.com/TheEhsanSarshar/d9852a7cc102dae0f07186e1dc4fe700
> 
> Hmm... just realised, that it might be worth it to have separate tests `toThenable`, then. There are a lot of ways in which this function can be used, and without tests it might be tricky to guarantee that all of them are handled correctly. Especially thinking about all the different ways in which errors could occur.

@barbogast are agree to finalize the api first and then write tests ?
barbogast: Sure, no worries. The only "hard" rule I suggest is having tests before merging the code. This might help during the implementation, so it could make sense to write tests next.
barbogast: Although... I've still got the nagging feeling that this approach might not be worth the effort and the complexity. This feels like a partial re-implementation of Promises. And yeah, if we wanna use `catch()` that's necessary for the code within `then()` to be executed synchronously. But to me it's not trivial to implement that correctly, get the edge cases right, writing tests, and later, when using it, being aware of the (probably subtle) differences to real promises.

Compared to all that the `try... catch` solution that works with the existing `makeThenable()` seems quite appealing, even if it looks awkward and requires one line of code duplication.

The reason that might make it necessary is if `makeThenable` cannot be typed correctly, though.

And, don't get me wrong, I'm totally on-board with the `toThenable()` approach from the gist. Just wanted to share my thoughts.
dai-shi: I'm kind of enjoying this effort. If it works as expected, it might be able to apply to other projects, namely jotai.
TheEhsanSarshar: > Although... I've still got the nagging feeling that this approach might not be worth the effort and the complexity. This feels like a partial re-implementation of Promises. And yeah, if we wanna use `catch()` that's necessary for the code within `then()` to be executed synchronously. But to me it's not trivial to implement that correctly, get the edge cases right, writing tests, and later, when using it, being aware of the (probably subtle) differences to real promises.
> 
> Compared to all that the `try... catch` solution that works with the existing `makeThenable()` seems quite appealing, even if it looks awkward and requires one line of code duplication.
> 
> The reason that might make it necessary is if `makeThenable` cannot be typed correctly, though.
> 
> And, don't get me wrong, I'm totally on-board with the `toThenable()` approach from the gist. Just wanted to share my thoughts.

just to make it clear dear @barbogast. the `try catch` solution will only work if all three values (storage, deserialize, migrate) are `sync`.
if one of them becomes async then we will lose the `try catch` context.
TheEhsanSarshar: > Thanks guys for working on this.
> I thought too that tests for `toThenable` would be nice.
> But, we haven't agreed on the api yet.
> Let me put my proposal here:
> 
> ```js
> const toThenable = (fn) => (input) => {
>   try {
>     const result = fn(input)
>     if (result instanceof Promise) {
>       return result
>     }
>     return {
>       then(onFulfilled) { return toThenable(onFulfilled)(result) },
>       catch(onRejected) { return this },
>     }
>   } catch (e) {
>     return {
>       then(onFulfilled) { return this },
>       catch(onRejected) { return toThenable(onRejected)(e) },
>     }
>   }
> }
> ```

thank you @dai-shi. I agree with this API. it's working like a charm 
todo list:
- [x] finalize API
- [ ] cherry-pick test provided by @barbogast 
- [ ] add types
- [ ] merge

TheEhsanSarshar: Good job. everybody.
now we need to make the sync tests pass.
barbogast: Will have a look later or tomorrow.
barbogast: With https://github.com/TheEhsanSarshar/zustand/pull/2 the tests are all green 🎉.

The solution feels a bit like a hack though, maybe there is a better way. One way to also fix it could be to later only override the state that was set by middlewares if the state is still empty. Might be a breaking change, though.
dai-shi: @TheEhsanSarshar wrt. https://gist.github.com/TheEhsanSarshar/d9852a7cc102dae0f07186e1dc4fe700#gistcomment-3762845

How about this? [playground](https://www.typescriptlang.org/play?#code/ATCWDsBcFMCcDMCGBjawAqALa5ECMAbaAHgDVECBXaAPmAG8AoEESbcMmgCgHtwAxSgXigCRACYAuYFwBuFatPJVoASmABeOqWAAfDO3xFOq6VhxGSpGsxbJEkZJk69wAJWgAraMhhSZsNCIAM580pTgANbgPADu4OpawDr65riEVjSmBhYZnLYAvoy2yHzBkMAAHprA4NCxwAAKsDwAtqDBJOCUrXhw3Fywqlr0BarFpeDlwJA8aZY1xB7BQpAANMAAkuAADpSQ3PDg0lwQe5DS2+eJdMurek0t7Z1L0CsEBw-zeXcfWZp0U67faXYGQbLfYy-T5JJgsGawACeDFs8Mm00C7wqGmARyB11RLFA8ACb3uEHKiHAqB4JOabQ6ahR8PhgUglFg4GAmPuIRy6ShZL+hJARRZ3Og7M5zPFM3YrkEwlEEnUcNlIDZHK5s0h0AVQhEYmg4lUgyFFT50PG6oKaxFdgcTlcHm8vmNqvtrMlWrlHWAfKpyIA9EHgPxNgANACyAFFgPYubE0IhZDxQOJ-eBEQB+T3AW32sUsArxx2YGRMtVeqVcqsstg4fVKo0mmXqzXSth+gNZ4AhsOR2OlxPJ1PpzM5vMF9X2RyYZ1eHx+D3qjXeztzQwZBeu5dcNRTu3iouiwpAA)

Also, please define `toThenable` outside `persist` and we can define some of functions in advance.
TheEhsanSarshar: @dai-shi please have a final review on this.
dai-shi: @AnatoleLucet Can you carefully review this?
barbogast: Right now the code doesn't behave correctly if the `version` in localStorage doesn't match the `version` in the code _and_ there is no `migrate` function. There is also no test covering this. 

Is this case supposed to be handled correctly? Thinking about it, a `console.error` warning might also be helpful in this case.
TheEhsanSarshar: > Right now the code doesn't behave correctly if the `version` in localStorage doesn't match the `version` in the code _and_ there is no `migrate` function. There is also no test covering this.
> 
> Is this case supposed to be handled correctly? Thinking about it, a `console.error` warning might also be helpful in this case.
> 
> --- edit
> PR: [TheEhsanSarshar#4](https://github.com/TheEhsanSarshar/zustand/pull/4). Also, the code handling the migration can be simplified. Would do that later.

Yeah, as you said a single console.warn would be enough.
dai-shi: Here you go: [playground](https://www.typescriptlang.org/play?#code/KYDwDg9gTgLgBDAnmYcDKMCGNUF44QBGAVsAMYwBQA9NQssAM5wCWAdo2C1MACZyFEcRsBgZsqAGZQIAWzgAlYJgoAaYcGAAuGnQAWMGGEZbaAcxYw9AV0IA6MnOoARYJPaXgAG0QAVBrwubh44Pv4ogYReRNQAbGQALACcwAkArAAcCQBMAMyZhJhpSUUADBnAAIy51aUA7EmVSRlJjZUJuaWBSCiM1DwqMNQAbpWx1Oy8oHa8djCMAMQAMgktALQrSWmUoJCw9ChwAAqYsCyYXuI4ADy+cKA4bLzMV8DqANL3II-PcADWwEQEEkcDu+ABQJBvgAfHBcJQ4HAAD5wAAURxYZD+tw+sJRvgAlAjkWjUYwsDgtKCCXDYRisTi4O88dSduBoPAeqhXmhvOQYNBbl8fi8KW84ABVWH4MliqmE2mStl7TkMOAAUQAjtYLpZEABhPTkAFQW7StHkiTy9RsYAAd1e8ppuFhhAgEC8yjYyo5B25YqWLHJwFtpphcItctB6jAPGGLAg1kYjupiuGEBYvB9+y56DFaC8mOAgeDobNEbJhbI2mjcFjwHjieTVZrCpdcHTmez8HYOCgkhU3NsjDIUBYhGAQoeId+r1hAG9iajC6XgFAqa8S4812aCVTUc7YZ2s4jrjyW9DUcTESvt+u8xIC0WtyGd+ei9DVNeNJ6KNAAPwbvmfJ-mG6jvtWn7fsA2q6kgABibCARqsGFkghrGm++YXsSe5ooeHYZlmAC+lDdn66CiK8U7fDOooSOai6np805PMwELAqCl5gKcMDnF4VInGcFzUb4uLqDwYBeIOyFuh6Xp4celCkbsvq5gA4lRYo0SKD44OaB6Kr45G5q45IyEIMoEUpqn7L2a4DtWebQMAACCXA6XRenAAuxIiGIUa8gFEhmsSZhaVacCacFNwwn5w6juONZoAlY4TqFiJTOZECIFSZkwBZykmWqrz6gMAphsKXmvOo+pJgKshBa8EZNdpMIGX5ohUnV5JyK1Ehfoi4UwFS0WiVBiKYFwQEue5LChQRxmULmPIVZg4URkxcDDQAkjgsj7mwmCyDW5nsGYBFnWwZgkmw1heF4JJHDIshBpOV03Sid0PdCnUwHtwAHWiR0nVSH3qMMFzWKdBXndZRFPS9b3XMev2kSta3hQAalDk5oIxwiBQA3B2a6MAmbBUndsgTlAcDo2qRxk0GMAAPJgHxEAcGejHEtQABU-NwAAcsdqCcVYqC9VA62oKisj1QIqDWGwLDasANL89QxIgzDY7XXzgvXkLrlwJIKsUBTcA8DA1hQKr11wJghPQLLdjG6CRouzLG0K+SZuWHAAAGdqTBAdp2NEZAiZjwBBwA5MwU0sGi0BO2wTuMIgbBkKTUDk1zBAgpYBLu4icBC-BaegMdUmoJL3uy3Ajj3fwE7B65Wc5xgrvhUHZsvYoygUCL2AsMMwBl4i-MewAAlMA73fAhntlHMe98Axva0NWkb8hK+wqtG+GzP5dCxKIhO839VyBoY66gAXmuU8V57qA23btr8B9cChw9Svez4KwDODdpZu2NnPBemAl5wAAFJoFZsLOwH0WCSEQFvTq99CxP33paSkzkfbABxl4aGPNLqw0diiZ6chkYfV+oiAWp9p6Skvs7MgN95BZTXPxFgT8oAvyYa-WePEZbyHMm-Ru4VE7XygDwNg8BIYkMnpAtw0CvDwHgYguwIiRAYMykwbhj9gC4IKmDChF1FQ90IcQ0h+NEY0JEGeWONi8bQnoXARhHtnoNhDPARgchUCeFkMwaQt8JznUAbMPRAhpJYlvPvDiII0AEgANoAF0T4e1ZmwHwgCEBe1vEA2MEAUBnCYAI7ev89CeHifuRJ6BUkZIYUbM+cBTbmxzpzDOH97YRMwGwCAks6aohKV0i4NIOmWy5gIiRshMDsDNhbLpv8WD-3btHB6QCJySBcvkqWYprbAD0IgXgMsukzN8F7HpX9FmdKtn-R66yLien4JgSQfY9mEwkIc45pyx6FzTnaI0Gd+n3FkWnCAZB2GyL4BU4kXMlC-LOcAKxssTGRWSYqVEsoJDITQOoNcMgoDIXVOCqA8NMw0hREpZpgihY7RBKAiqQC8HAGkRPfOVtXqMDmTAMgehPlcyligMgqCWBAKNDwdQTKN4rP-gM+A7ckywo9pcoMrBmDKvNo9IFIYna8F4H0gQAw-gRP5f0jaAo4BAjtoAuFiIOUFyQlTawNM1yZNae0pZVsrWlJ2VAeQpTybBm-gc16ZgzkUwudUkJ3rC4PIARsl5VS9VBpZiyg5jqKbMG5by-lKyrCCttMIEVYqJVrmUWfSp4bkX7zTSGlM-TEAQ2ZlzF1bryWWPsdy1xRVKD2X7ION+R0oiThcb5REks2DXCxpeLm8F7ruE2bwfcijoZUhcQRLG3bkazpJJckMmBR0zuhHhA9I7PQnuJNHPlegT0jLYEoUgFA+D7gGP4ymGoyVbp3Y4vd+JgVHsvbOs9gHj2zr7Y4DgnIIDnqA3gOA1wlCMCXuoHabAwDWBgJeYkkhP2onYJhkacB0NEYIshmBVCkaOIo+olkcHj20ew5QAiBGMNYapKRrDoHD2MaYEvaUE6EBQCEFtREUH-Y8BQ+oiMeG2Nke-KgtEUmYHsEtDnYAnFqE9ppGJ8u1zDnSfgJgZgDHL1MfcYiUi5dDm23tnAPTk7gUPoXV4JdLzdPfhswZgUZngAucXasjzqIVMyZM0PIzRIbNWcGtFm9-KH1Pv5HwTz0X9OiE-vk9V4W-PXCbZZ8uxFYtWeJMRZu2B82og1g578BnHN7LYAFtzQWUs1bS4iHzMbM7Dvg3ltgiACsxa8+V29iXgDPpwLwVL7WfOwbA56MbE2UtVai9For35rP0z7bZeAEn4D1vgPgM8VU2LeRw+Jrm7gzBAQkGVZQFUeaxdGdmwSzNyTsy6YwHmLGsX+Q3BFG4+N1DDVGgD1x6hk4zR4HNMhG5FRbT2218uutivbV3oQisBE16XFjqjkQWDeEIc0UglBaDus4vwai7GuMyGKnBt+LhBOn4RmJ9o04rDmCynvB9AiVOiE06B9+KIKg-i3lR0CmpLNUcIqOSc5FfPUdZsLvgUoqOa0DVKxGZ7XNmBIhRPOUixJPR+NjjdnAfOSQqwXuwPgZFJ0iaR4iMBG18DDT5weTX8WBUrcdx4ug9YJ7yLBUS1gjKvbO8CcwBVcBrc3Kq3YMwdgU0Z3xxy4QmZ34zm4Y7Z2PFwqrcN4iJTqIACEEfpsdYy-ZqDV2rxpdRHYRvpwzCMAIvVi7HB5J2DtKcRrw2bNByZpy-2r0DWeh7zwKkKt4MIAgHAawYBeDfIACTzl1kVz5FhA+SMjzIuRMBcnT8hqs+Ddgg797gKttL-kG9N6gC3q-hXUc7xgM-p2XBvwF7t83HXnJ5sosMWwQQ183-zJEAMJyJGvV-w0H+n2kx3lH-xRiIkE190R1ZU2jgEbwT1EEMkL3LmLwl1CBZgr3LiJlRFZhIH5DsAhEYAp2q3C1RHqWSXSVLn9XVBUD0GxQhDbwvxL0IO8BZjsHYDIBISykYMBBpAADJJCY8+QcAvkcAUkIQmk1tH9Nsi8QRURhc4liDfdERtDRdBC2COCuCJDFQphjd9kJAlDAQ0k1Dv9K87MQF-9eQmd-N5wFDxQlcM5iJS4p1sV8ceEn5eBN14dhsDNZRZVwsj5CFS5-IAZZBUQUc74gi+BN0Nsv8oDoNhBMAJ5eB+p5D8Bk5kEwdv9ij-JmoZRWUJJgApJBweCbNGBci+ACj-NqjDk6jqxH9jwYCEiPcSs+Y6AeAkVvkQh+JPCC0BVQAWYIlpYM0JBv9aAr47RoA-hThEwnhZ9hAPQJ5PkgwUNxYQRo9pYIlhi5cJBQ1RiU9s5c4I9BjPkb9WUaQ11UAVjW4AFpIPkIAOUxwDU9VQ5C0xiLgJjBBm5yp-NVslirVciEYrAgx1A7RUATp7965w8DlQl5BscLdew58G4gTHpWUp5jdPD4IXo+d0M0Abi4cUQrdggv4sj-ZIByREVziul9RnlCgsQtdH1Zc-lzdY5-w7BURhoDwqUaSng6TbcGE6BCBJgNV58RB+BoSeig5whgBSUiVONNkzBgT2B0wb0KZ+5CVoBiQQDeMFsI9sDYDAY7BZSnhIjYiCQkixZH85hnMHTZZQj2x6ti8I9Qj29bNMtGdUiPTqclFH8SsbNfDvw3SQxsVgyjF8jnFcYGjoti8EygCkyN5-SL90yDE3CszrFcY7BvC4AS9cB8BvCSC0ti91ccBqyZsq8M46z-ML80sMzCdCzPTizqi2zosOzgi+cXESzW1vR2tosIy1s+yJMu9jSoA69xzy4g5mpohMAph+AMSd8f9W42B45FUkSWAI0Lj08NM4ABk4AWzbkpkM4e9mBil4x1zz9FzL8L8ytvBL4AyncxRSS+pY4KSbiIwByWjkylFkExQ+yb8gKuywzoYwKJBJzIypyozXT-D5ZDzkUkz4KwiazNCWzML6y9Dy5WUfzGo-y2BKSc4Iw8LmpoiL8b9qKDlwshI+IY4Qp8V-hAROJ8YELAz7N4j9p+ikLCsUL3TUybMmSYAWS+SKZ2SHpOS-hBThScCCR1BaT3Av41DXSvdsUawNToAxLy4JKpLI0uZZKvB5LFK1KbdeACVNLv9Eca9DymMIxHKzAFzERb87Bm9W9sLosb8sDvLJyej+LAZBKhsbNhpUdk5cJv8DNiKySyKKKyBvx-wHNMDG9XKmN1AsD4rfyN5-zKL1C4AqRMr+N1FlIgA)
TheEhsanSarshar: Good job everybody. 
Thank you @dai-shi @barbogast @AnatoleLucet for working on this 
barbogast: Like-wise, after years working with open-source software I'm still amazed getting such high-quality software for free. And in addition to that it feels like there's 24/7 support with direct access to the core-developers, no premium-support-subscription needed.
But it's not one-sided since, as seen here, users are happy to contribute and collaborate, and so it's a common effort, rather than provider-consumer. I feel that the open source development is one of the best and highest value-generating way of collaboration that humans have.

Hehe, got all philosophical here 🙈 Anyway, thanks all 😍 

Patch:
diff --git a/.size-snapshot.json b/.size-snapshot.json
index 17491b4..41cf8b2 100644
--- a/.size-snapshot.json
+++ b/.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,
diff --git a/src/middleware.ts b/src/middleware.ts
index 7530a48..3097a52 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -204,6 +204,43 @@ 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>
@@ -211,8 +248,8 @@ export const persist = <S extends State>(
   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,38 +314,52 @@ 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()
@@ -303,4 +367,8 @@ export const persist = <S extends State>(
     get,
     api
   )
+
+  return stateFromStorageInSync
+    ? { ...configResult, ...stateFromStorageInSync }
+    : configResult
 }


Test patch:
diff --git a/tests/persist.test.tsx b/tests/persist.test.tsx
index 96b5002..a5a105a 100644
--- a/tests/persist.test.tsx
+++ b/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)
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
new file mode 100644
index 0000000..e8539d9
--- /dev/null
+++ b/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)
+  })
+})