fix: flushPending in async write (#2804)

* fix: flushPending in async write * isSync in onMount to * apply the change * add test for flushPending (#2809) * Fix/flush pending finally everywhere (#2816) * call onmount flushPending in finally block * temp remove unmount test cases. to be fixed later with pending injection (#2810) * flushpending everywhere * Apply suggestions from code review --------- Co-authored-by: Daishi Kato <[email protected]> * fix type assertion --------- Co-authored-by: David Maskasky <[email protected]>
pmndrs · Nov 14, 2024 · 6ff4eb2 · 6ff4eb2
1 parent 47fef4d
commit 6ff4eb2
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 35 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -333,15 +333,18 @@ const buildStore = (
       }
       // a !== atom
       const aState = readAtomState(pending, a, dirtyAtoms)
-      if (isSync) {
-        addDependency(pending, atom, atomState, a, aState)
-      } else {
-        const pending = createPending()
-        addDependency(pending, atom, atomState, a, aState)
-        mountDependencies(pending, atom, atomState)
-        flushPending(pending)
+      try {
+        return returnAtomValue(aState)
+      } finally {
+        if (isSync) {
+          addDependency(pending, atom, atomState, a, aState)
+        } else {
+          const pending = createPending()
+          addDependency(pending, atom, atomState, a, aState)
+          mountDependencies(pending, atom, atomState)
+          flushPending(pending)
+        }
       }
-      return returnAtomValue(aState)
     }
     let controller: AbortController | undefined
     let setSelf: ((...args: unknown[]) => unknown) | undefined
@@ -485,46 +488,56 @@ const buildStore = (
     atom: WritableAtom<Value, Args, Result>,
     ...args: Args
   ): Result => {
+    let isSync = true
     const getter: Getter = <V>(a: Atom<V>) =>
       returnAtomValue(readAtomState(pending, a))
     const setter: Setter = <V, As extends unknown[], R>(
       a: WritableAtom<V, As, R>,
       ...args: As
     ) => {
       const aState = getAtomState(a)
-      let r: R | undefined
-      if (isSelfAtom(atom, a)) {
-        if (!hasInitialValue(a)) {
-          // NOTE technically possible but restricted as it may cause bugs
-          throw new Error('atom not writable')
+      try {
+        if (isSelfAtom(atom, a)) {
+          if (!hasInitialValue(a)) {
+            // NOTE technically possible but restricted as it may cause bugs
+            throw new Error('atom not writable')
+          }
+          const hasPrevValue = 'v' in aState
+          const prevValue = aState.v
+          const v = args[0] as V
+          setAtomStateValueOrPromise(a, aState, v)
+          mountDependencies(pending, a, aState)
+          if (!hasPrevValue || !Object.is(prevValue, aState.v)) {
+            addPendingAtom(pending, a, aState)
+            recomputeDependents(pending, a, aState)
+          }
+          return undefined as R
+        } else {
+          return writeAtomState(pending, a, ...args)
         }
-        const hasPrevValue = 'v' in aState
-        const prevValue = aState.v
-        const v = args[0] as V
-        setAtomStateValueOrPromise(a, aState, v)
-        mountDependencies(pending, a, aState)
-        if (!hasPrevValue || !Object.is(prevValue, aState.v)) {
-          addPendingAtom(pending, a, aState)
-          recomputeDependents(pending, a, aState)
+      } finally {
+        if (!isSync) {
+          flushPending(pending)
         }
-      } else {
-        r = writeAtomState(pending, a, ...args) as R
       }
-      flushPending(pending)
-      return r as R
     }
-    const result = atomWrite(atom, getter, setter, ...args)
-    return result
+    try {
+      return atomWrite(atom, getter, setter, ...args)
+    } finally {
+      isSync = false
+    }
   }
 
   const writeAtom = <Value, Args extends unknown[], Result>(
     atom: WritableAtom<Value, Args, Result>,
     ...args: Args
   ): Result => {
     const pending = createPending()
-    const result = writeAtomState(pending, atom, ...args)
-    flushPending(pending)
-    return result
+    try {
+      return writeAtomState(pending, atom, ...args)
+    } finally {
+      flushPending(pending)
+    }
   }
 
   const mountDependencies = (
@@ -575,11 +588,29 @@ const buildStore = (
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
         addPendingFunction(pending, () => {
-          const onUnmount = atomOnMount(atom, (...args) =>
-            writeAtomState(pending, atom, ...args),
-          )
-          if (onUnmount) {
-            mounted.u = onUnmount
+          let isSync = true
+          try {
+            const onUnmount = atomOnMount(atom, (...args) => {
+              try {
+                return writeAtomState(pending, atom, ...args)
+              } finally {
+                if (!isSync) {
+                  flushPending(pending)
+                }
+              }
+            })
+            if (onUnmount) {
+              mounted.u = () => {
+                isSync = true
+                try {
+                  onUnmount()
+                } finally {
+                  isSync = false
+                }
+              }
+            }
+          } finally {
+            isSync = false
           }
         })
       }

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
@@ -579,3 +579,70 @@ it('should update derived atom even if dependances changed (#2697)', () => {
   store.set(primitiveAtom, 1)
   expect(onChangeDerived).toHaveBeenCalledTimes(1)
 })
+
+describe('should invoke flushPending only after all atoms are updated (#2804)', () => {
+  const store = createStore()
+
+  it('should invoke flushPending only after all atoms are updated with set', () => {
+    const a = atom(0)
+    const setResult = []
+    const w = atom(null, (_get, set, value: number) => {
+      setResult.push('before set')
+      set(a, value)
+      setResult.push('after set')
+    })
+    store.sub(a, () => {
+      setResult.push('a value changed - ' + store.get(a))
+    })
+    setResult.push('before store.set')
+    store.set(w, 1)
+    setResult.push('after store.set')
+    expect(setResult).not.toEqual([
+      'before store.set',
+      'before set',
+      'a value changed - 1',
+      'after set',
+      'after store.set',
+    ])
+    expect(setResult).toEqual([
+      'before store.set',
+      'before set',
+      'after set',
+      'a value changed - 1',
+      'after store.set',
+    ])
+  })
+
+  it('should invoke flushPending only after all atoms are updated with mount', () => {
+    const mountResult = []
+    const a = atom(0)
+    const m = atom(null, (_get, set, value: number) => {
+      set(a, value)
+    })
+    m.onMount = (setSelf) => {
+      mountResult.push('before onMount setSelf')
+      setSelf(1)
+      mountResult.push('after onMount setSelf')
+    }
+    mountResult.push('before store.sub')
+    store.sub(a, () => {
+      mountResult.push('a value changed - ' + store.get(a))
+    })
+    const unsub = store.sub(m, () => {})
+    mountResult.push('after store.sub')
+    expect(mountResult).not.toEqual([
+      'before store.sub',
+      'before onMount setSelf',
+      'a value changed - 1',
+      'after onMount setSelf',
+      'after store.sub',
+    ])
+    expect(mountResult).toEqual([
+      'before store.sub',
+      'before onMount setSelf',
+      'after onMount setSelf',
+      'a value changed - 1',
+      'after store.sub',
+    ])
+  })
+})