diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index c7390cb45d..457a64b245 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -595,7 +595,14 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { return result }) if (onUnmount) { - mounted.u = 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 index c0be322c45..92401f2431 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -579,3 +579,86 @@ 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') + return () => { + mountResult.push('before onUnmount setSelf') + setSelf(2) + mountResult.push('after onUnmount 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') + mountResult.push('before store.unsub') + unsub() + mountResult.push('after store.unsub') + expect(mountResult).not.toEqual([ + 'before store.sub', + 'before onMount setSelf', + 'a value changed - 1', + 'after onMount setSelf', + 'after store.sub', + 'before store.unsub', + 'before onUnmount setSelf', + 'a value changed - 2', + 'after onUnmount setSelf', + 'after store.unsub', + ]) + expect(mountResult).toEqual([ + 'before store.sub', + 'before onMount setSelf', + 'after onMount setSelf', + 'a value changed - 1', + 'after store.sub', + 'before store.unsub', + 'before onUnmount setSelf', + 'after onUnmount setSelf', + 'a value changed - 2', + 'after store.unsub', + ]) + }) +})