diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 5b0c47cf17..190d0766de 100644 --- 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,6 +488,7 @@ const buildStore = ( atom: WritableAtom, ...args: Args ): Result => { + let isSync = true const getter: Getter = (a: Atom) => returnAtomValue(readAtomState(pending, a)) const setter: Setter = ( @@ -492,29 +496,36 @@ const buildStore = ( ...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 = ( @@ -522,9 +533,11 @@ const buildStore = ( ...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 index 4fa0e668af..859af8730c 100644 --- 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', + ]) + }) +})