fix: setAtom uses stale pending on atom unmount (#2811)

* fix: setAtom should use the correct pending on unmount * ideating with overwritable setAtom pattern * encapsulate isSync inside setAtom * refactor * refactor 2 * refactor to use createRuntime * refactor 3 * how about this? --------- Co-authored-by: daishi <daishi@axlight.com>
pmndrs · Nov 17, 2024 · 29662ea · 29662ea
1 parent 0f05df8
commit 29662ea
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 21 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -77,7 +77,7 @@ type Mounted = {
   /** Set of mounted atoms that depends on the atom. */
   readonly t: Set<AnyAtom>
   /** Function to run when the atom is unmounted. */
-  u?: OnUnmount
+  u?: (pending: Pending) => void
 }
 
 /**
@@ -602,31 +602,31 @@ const buildStore = (
       }
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
-        addPendingFunction(pending, () => {
+        let setAtom: (...args: unknown[]) => unknown
+        const createInvocationContext = <T>(pending: Pending, fn: () => T) => {
           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
-                }
+          setAtom = (...args: unknown[]) => {
+            try {
+              return writeAtomState(pending, atom, ...args)
+            } finally {
+              if (!isSync) {
+                flushPending(pending)
               }
             }
+          }
+          try {
+            return fn()
           } finally {
             isSync = false
           }
+        }
+        addPendingFunction(pending, () => {
+          const onUnmount = createInvocationContext(pending, () =>
+            atomOnMount(atom, (...args) => setAtom(...args)),
+          )
+          if (onUnmount) {
+            mounted.u = (pending) => createInvocationContext(pending, onUnmount)
+          }
         })
       }
     }
@@ -646,7 +646,7 @@ const buildStore = (
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        addPendingFunction(pending, onUnmount)
+        addPendingFunction(pending, () => onUnmount(pending))
       }
       delete atomState.m
       if (import.meta.env?.MODE !== 'production') {

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
@@ -896,3 +896,17 @@ it('throws falsy errors in onMount, onUnmount, and listeners', () => {
   })
   expect(() => store.set(c, 1)).toThrow('')
 })
+
+it('should use the correct pending on unmount', () => {
+  const store = createStore()
+  const a = atom(0)
+  const b = atom(0, (_, set, update: number) => set(a, update))
+  b.onMount = (setAtom) => () => setAtom(1)
+  const aListener = vi.fn()
+  store.sub(a, aListener)
+  const unsub = store.sub(b, () => {})
+  aListener.mockClear()
+  unsub()
+  expect(store.get(a)).toBe(1)
+  expect(aListener).toHaveBeenCalledTimes(1)
+})