fix(vanilla): store should flush pending write triggered asynchronous…

…ly and indirectly (#2451)

* fix(vanilla): store should flush pending write triggered asynchronously and indirectly

* add PR number to the test case

* an alternative way to flush async writes

src/vanilla/store.ts

@@ -560,12 +560,15 @@ export const createStore = () => {
     atom: WritableAtom<Value, Args, Result>,
     ...args: Args
   ): Result => {
-    let isSync = true
     const getter: Getter = <V>(a: Atom<V>) => returnAtomValue(readAtomState(a))
     const setter: Setter = <V, As extends unknown[], R>(
       a: WritableAtom<V, As, R>,
       ...args: As
     ) => {
+      const isSync = pendingStack.length > 0
+      if (!isSync) {
+        pendingStack.push(new Set([a]))
+      }
       let r: R | undefined
       if (isSelfAtom(atom, a)) {
         if (!hasInitialValue(a)) {
@@ -581,7 +584,7 @@ export const createStore = () => {
         r = writeAtomState(a as AnyWritableAtom, ...args) as R
       }
       if (!isSync) {
-        const flushed = flushPending([a])
+        const flushed = flushPending(pendingStack.pop()!)
         if (import.meta.env?.MODE !== 'production') {
           storeListenersRev2.forEach((l) =>
             l({ type: 'async-write', flushed: flushed! }),
@@ -591,7 +594,6 @@ export const createStore = () => {
       return r as R
     }
     const result = atom.write(getter, setter, ...args)
-    isSync = false
     return result
   }
 

tests/vanilla/store.test.tsx

@@ -502,3 +502,29 @@ it('should mount once with atom creator atom (#2314)', async () => {
   store.sub(atomCreatorAtom, () => {})
   expect(countAtom.onMount).toHaveBeenCalledTimes(1)
 })
+
+it('should flush pending write triggered asynchronously and indirectly (#2451)', async () => {
+  const store = createStore()
+  const anAtom = atom('initial')
+
+  const callbackFn = vi.fn((_value: string) => {})
+  const unsub = store.sub(anAtom, () => {
+    callbackFn(store.get(anAtom))
+  })
+
+  const actionAtom = atom(null, async (_get, set) => {
+    await Promise.resolve() // waiting a microtask
+    set(indirectSetAtom)
+  })
+
+  const indirectSetAtom = atom(null, (_get, set) => {
+    set(anAtom, 'next')
+  })
+
+  // executing the chain reaction
+  await store.set(actionAtom)
+
+  expect(callbackFn).toHaveBeenCalledOnce()
+  expect(callbackFn).toHaveBeenCalledWith('next')
+  unsub()
+})