diff --git a/src/utils/__tests__/batched-updates.test.js b/src/utils/__tests__/batched-updates.test.js index eeade58..82896a4 100644 --- a/src/utils/__tests__/batched-updates.test.js +++ b/src/utils/__tests__/batched-updates.test.js @@ -56,13 +56,20 @@ describe('batch', () => { const child = jest.fn().mockReturnValue(null); render({child}); const update = child.mock.calls[0][0]; - act(() => update()); + act(() => { + update(); + update(); + update(); + }); + + // nothing should be yet called + expect(child.mock.calls[2]).toEqual(undefined); // scheduler uses timeouts on non-browser envs await act(() => new Promise((r) => setTimeout(r, 10))); // assertion no longer relevant with React 18+ - expect(child.mock.calls[2]).toEqual([expect.any(Function), 1, 1]); + expect(child.mock.calls[2]).toEqual([expect.any(Function), 3, 1]); supportsMock.mockRestore(); }); diff --git a/src/utils/batched-updates.js b/src/utils/batched-updates.js index 5fcf8b8..6ac286a 100644 --- a/src/utils/batched-updates.js +++ b/src/utils/batched-updates.js @@ -8,23 +8,33 @@ import { import defaults from '../defaults'; import supports from './supported-features'; -let isInsideBatchedSchedule = false; +let batchedScheduled = false; + +let batched = []; + +const executeBatched = () => { + unstable_batchedUpdates(() => { + while (batched.length) { + const currentBatched = batched; + batched = []; + currentBatched.forEach((fn) => fn()); + } + // important to reset it before exiting this function + // as React will dump everything right after + batchedScheduled = false; + }); +}; export function batch(fn) { // if we are in node/tests or nested schedule - if ( - !defaults.batchUpdates || - !supports.scheduling() || - isInsideBatchedSchedule - ) { + if (!defaults.batchUpdates || !supports.scheduling()) { return unstable_batchedUpdates(fn); } - isInsideBatchedSchedule = true; - // Use ImmediatePriority as it has -1ms timeout - // https://github.com/facebook/react/blob/main/packages/scheduler/src/forks/Scheduler.js#L65 - return scheduleCallback(ImmediatePriority, function scheduleBatchedUpdates() { - unstable_batchedUpdates(fn); - isInsideBatchedSchedule = false; - }); + batched.push(fn); + + if (!batchedScheduled) { + batchedScheduled = true; + return scheduleCallback(ImmediatePriority, executeBatched); + } }