Infinite re-render with async atomWithStorage and React 19

[friggeri]

Bug Description

I'm using the latest version of jotai alongside React 19. I have the following atom defined:

const counterAtom = atomWithStorage('counter', 0, {
  async getItem(_key: string, _initialValue: number) {
    console.log('getting item');
    return await Promise.resolve(storedValue);
  },
  async setItem(_key, newValue) {
    storedValue = await new Promise((resolve) => resolve(newValue));
  },
  async removeItem() {},
});

In absence of any calls to setting the atom, my expectation would be that this would only trigger two renders: one before the promise is resolved, rendering the suspense fallback, and once once the promise is resolved, rendering the tree. Instead I'm getting an infinite loop of going back and forth between fallback and suspense tree.

Repro showcasing the issue in 19.0.0: https://stackblitz.com/edit/vitejs-vite-z4b5kyy1?file=src%2FApp.tsx
This worked in React 18.3.1: https://stackblitz.com/edit/vitejs-vite-vhskxk28?file=package.json,src%2FApp.tsx

Reproduction Link

https://stackblitz.com/edit/vitejs-vite-z4b5kyy1?file=src%2FApp.tsx

[dai-shi]

Thanks for reporting. We should investigate this.

To isolate the problem, what we want first is:

1. reproduce the behavior with atom without atomWithStorage
2. reproduce the behavior with use without Jotai

I would try when I have time, but if anyone is interested, feel free!

[friggeri]

@dai-shi here is a smaller repro without using atomWithStorage, the issue is a combination of using onMount to set a promise, and strict mode, if I disable the later it works:
https://stackblitz.com/edit/vitejs-vite-rrg32jrm?file=src%2FApp.tsx

[dai-shi]

Thanks! This looks like the new Suspense behavior introduced in 19. I'll see if I can do anything. Needs time for it. Unfortunately, I can't think of any workaround with atomWithStorage case.

[dai-shi]

Opened a PR for a failing test: #2933

[dai-shi]

After looking into it, this behavior is expected as React remounts the atom. So, { delay: 0 } would be the actual solution. Now, the real issue is with atomWithStorage. Let me see.

[eastdraconis]

const [counter, setCounter] = useAtom(counterAtom,{delay:0});
I tested it and if you give a delay option to useAtom, it works normally
I think internally useAtomValue is calling reender via setTimeout
Is it related to this?

[dai-shi]

Oh, yeah. That's the workaround. (I think I tried it, but doing it now again shows working good.)

It is related, but it might not be an ideal fix.

[dai-shi]

Returning a stable promise seems like a canonical solution.

https://github.com/pmndrs/jotai/pull/2933/files#diff-b947755038dbc008858a31a1209b392de371aa68cdb77952278c174f282bd39bR386-R391

[dai-shi]

As StrictMode in React 19 remounts atoms, we need to return stable promises.

Would it possible for you to do something like this?
https://github.com/pmndrs/jotai/pull/2933/files#diff-e6f199768dd7a3a78f35b26f83ed1c2101d98d71b10de53b2c97a77874b0b8f3R751-R768

    let storedValue = 0
    let cachedPromise:
      | [typeof storedValue, Promise<typeof storedValue>]
      | null = null
    const counterAtom = atomWithStorage('counter', 0, {
      getItem(_key: string, _initialValue: number) {
        if (cachedPromise && cachedPromise[0] === storedValue) {
          return cachedPromise[1]
        }
        const promise = Promise.resolve(storedValue)
        cachedPromise = [storedValue, promise]
        return promise
      },
      async setItem(_key, newValue) {
        storedValue = await new Promise((resolve) => resolve(newValue))
      },
      async removeItem() {},
    })

[eastdraconis]

Is there no problem if I don't initialize cahcedPromise at setItem?

[dai-shi]

As it's just a cache, setItem doesn't need anything.

[dhruvkelawala]

Seems like its also related to #2848