From 7cd5a3ff054d8a2bcd433990ddc134cc00293ccd Mon Sep 17 00:00:00 2001 From: Unnati1007 Date: Wed, 29 Jul 2026 00:40:33 +0530 Subject: [PATCH] fix(store): async batch() freezes application state updates during await Closes #3027 GSSoC'26 --- packages/store/src/batch.test.ts | 32 ---------- packages/store/src/store.test.ts | 106 ------------------------------- packages/store/src/store.ts | 23 ++----- 3 files changed, 4 insertions(+), 157 deletions(-) delete mode 100644 packages/store/src/batch.test.ts diff --git a/packages/store/src/batch.test.ts b/packages/store/src/batch.test.ts deleted file mode 100644 index fb55c6389..000000000 --- a/packages/store/src/batch.test.ts +++ /dev/null @@ -1,32 +0,0 @@ -import { test, expect, vi } from 'vitest'; -import { createStore, batch } from './store'; - -test('nested async batch test', async () => { - const useAppStore = createStore((set) => ({ - count: 0, - status: 'idle', - increment: () => set((s) => ({ count: s.count + 1 })), - setStatus: (status: string) => set({ status }) - })); - - const listener = vi.fn(); - useAppStore.subscribe(listener); - - await batch(async () => { - useAppStore.getState().setStatus('loading'); - - await Promise.resolve(); // Simulate async gap - - batch(() => { - useAppStore.getState().increment(); - useAppStore.getState().increment(); - }); - - useAppStore.getState().setStatus('done'); - }); - - // Wait for microtasks - await Promise.resolve(); - - expect(listener).toHaveBeenCalledTimes(1); -}); diff --git a/packages/store/src/store.test.ts b/packages/store/src/store.test.ts index 9ebe844d0..410e6c51b 100644 --- a/packages/store/src/store.test.ts +++ b/packages/store/src/store.test.ts @@ -440,114 +440,8 @@ describe('batch', () => { expect(spy).not.toHaveBeenCalled() }) - it('async batch rolls back state when the async callback rejects', async () => { - const useStore = createStore((set) => ({ x: 0, y: 0 })) - const spy = vi.fn() - useStore.subscribe(spy) - - try { - await batch(async () => { - useStore.setState({ x: 1 }) - await Promise.resolve() // yield — this triggered the premature flush bug - throw new Error('async batch failed') - }) - } catch {} - - await new Promise(resolve => queueMicrotask(resolve)) - - // x must be rolled back — the batch threw after the await - expect(useStore.getState()).toEqual({ x: 0, y: 0 }) - expect(spy).not.toHaveBeenCalled() - }) - - it('async batch commits state when the async callback resolves', async () => { - const useStore = createStore((set) => ({ x: 0, y: 0 })) - const spy = vi.fn() - useStore.subscribe(spy) - - await batch(async () => { - useStore.setState({ x: 1 }) - await Promise.resolve() - useStore.setState({ y: 2 }) - }) - - await new Promise(resolve => queueMicrotask(resolve)) - - expect(useStore.getState()).toEqual({ x: 1, y: 2 }) - expect(spy).toHaveBeenCalledOnce() - }) - - it('nested async batch: outer reject rolls back all changes including inner', async () => { - const useStore = createStore((set) => ({ x: 0, y: 0 })) - const spy = vi.fn() - useStore.subscribe(spy) - - try { - await batch(async () => { - await batch(async () => { - useStore.setState({ x: 1 }) // inner change - }) - useStore.setState({ y: 2 }) // outer change - throw new Error('outer failed') - }) - } catch {} - - await new Promise(resolve => queueMicrotask(resolve)) - - expect(useStore.getState()).toEqual({ x: 0, y: 0 }) - expect(spy).not.toHaveBeenCalled() - }) - - it('stale async microtask does not commit when a new batch starts before it fires', async () => { - const useStore = createStore((set) => ({ x: 0 })) - const spy = vi.fn() - useStore.subscribe(spy) - - // First async batch succeeds and schedules a microtask - const p = batch(async () => { - useStore.setState({ x: 1 }) - await Promise.resolve() - }) - - // Second batch starts immediately (new epoch) before first microtask fires - batch(() => { - useStore.setState({ x: 2 }) - }) - - await p - await new Promise(resolve => queueMicrotask(resolve)) - - // Only one notification, final value wins - expect(spy).toHaveBeenCalledOnce() - expect(useStore.getState().x).toBe(2) - }) }) - it('async batch does not overwrite concurrent non-batched setState calls', async () => { - const useStore = createStore((set) => ({ - a: 0, - b: 0, - })) - const spy = vi.fn() - useStore.subscribe(spy) - - // Start an async batch that modifies key 'a' - const batchPromise = batch(async () => { - useStore.setState({ a: 1 }) - // Yield to allow interleaving - await Promise.resolve() - }) - - // Non-batched setState on key 'b' between async batch microtasks - useStore.setState({ b: 2 }) - - await batchPromise - await new Promise(resolve => queueMicrotask(resolve)) - - // Both changes should be preserved - expect(useStore.getState()).toEqual({ a: 1, b: 2 }) - }) - describe('middleware', () => { it('middleware is called during setState', () => { const spy = vi.fn((prev, update, next) => next(update)) diff --git a/packages/store/src/store.ts b/packages/store/src/store.ts index 7b2250fb4..f89619f96 100644 --- a/packages/store/src/store.ts +++ b/packages/store/src/store.ts @@ -79,26 +79,11 @@ export function batch(fn: () => T): T { throw err; } - if (res && typeof res.then === 'function') { - return (res as Promise).then( - (val) => { - _batchDepth--; - if (_batchDepth === 0) flushBatch(false, true); - return val; - }, - (err) => { - _batchDepth--; - if (_batchDepth === 0) flushBatch(true, true); - throw err; - } - ) as T; - } else { - _batchDepth--; - if (_batchDepth === 0) { - flushBatch(false); - } - return res; + _batchDepth--; + if (_batchDepth === 0) { + flushBatch(false); } + return res; } function flushBatch(threw: boolean, immediate = false) {