Fix nested batch scheduling (#120)

babel.config.js

@@ -5,9 +5,15 @@ module.exports = {
     [
       '@babel/preset-env',
       {
-        targets: { edge: '16' },
+        targets: [
+          'last 2 chrome versions',
+          'last 2 firefox versions',
+          'last 2 safari versions',
+          'last 2 and_chr versions',
+          'last 2 ios_saf versions',
+          'edge >= 18',
+        ],
         modules: false,
-        exclude: ['transform-typeof-symbol'],
         loose: true,
       },
     ],

src/index.js

@@ -5,5 +5,5 @@ export { createSubscriber } from './components/subscriber';
 export { createHook } from './components/hook';
 export { default as defaults } from './defaults';
 export { createStore, defaultRegistry } from './store';
-export { unstable_batchedUpdates as batch } from './utils/batched-updates';
+export { batch } from './utils/batched-updates';
 export { createSelector } from './utils/create-selector';

src/utils/__tests__/batched-updates.test.js

@@ -0,0 +1,70 @@
+/* eslint-env jest */
+
+import React, { useState } from 'react';
+import { mount } from 'enzyme';
+
+import { createHook } from '../../components/hook';
+import defaults from '../../defaults';
+import { createStore, defaultRegistry } from '../../store';
+import supports from '../../utils/supported-features';
+import { batch } from '../batched-updates';
+
+const Store = createStore({
+  initialState: { count: 0 },
+  actions: {
+    increment: () => ({ getState, setState }) => {
+      setState({ count: getState().count + 1 });
+    },
+  },
+});
+
+const useHook = createHook(Store);
+
+describe('batch', () => {
+  const TestComponent = ({ children }) => {
+    const [{ count }, actions] = useHook();
+    const [localCount, setLocalCount] = useState(0);
+    const update = () =>
+      batch(() => {
+        actions.increment();
+        setLocalCount(localCount + 1);
+      });
+
+    return children(update, count, localCount);
+  };
+
+  beforeEach(() => {
+    defaultRegistry.stores.clear();
+  });
+
+  it('should batch updates with scheduling disabled', () => {
+    const child = jest.fn().mockReturnValue(null);
+    mount(<TestComponent>{child}</TestComponent>);
+    const update = child.mock.calls[0][0];
+    update();
+
+    expect(child.mock.calls).toHaveLength(2);
+    expect(child.mock.calls[1]).toEqual([expect.any(Function), 1, 1]);
+  });
+
+  it('should batch updates with scheduling enabled', async () => {
+    const supportsMock = jest
+      .spyOn(supports, 'scheduling')
+      .mockReturnValue(true);
+    defaults.batchUpdates = true;
+
+    const child = jest.fn().mockReturnValue(null);
+    mount(<TestComponent>{child}</TestComponent>);
+    const update = child.mock.calls[0][0];
+    update();
+
+    // scheduler uses timeouts on non-browser envs
+    await new Promise((r) => setTimeout(r, 10));
+
+    expect(child.mock.calls).toHaveLength(2);
+    expect(child.mock.calls[1]).toEqual([expect.any(Function), 1, 1]);
+
+    supportsMock.mockRestore();
+    defaults.batchUpdates = false;
+  });
+});

src/utils/batched-updates.js

@@ -1,2 +1,33 @@
 /* eslint-disable import/no-unresolved */
-export { unstable_batchedUpdates } from 'react-dom';
+import { unstable_batchedUpdates } from 'react-dom';
+import {
+  unstable_scheduleCallback as scheduleCallback,
+  unstable_UserBlockingPriority as UserBlockingPriority,
+} from 'scheduler';
+
+import defaults from '../defaults';
+import supports from './supported-features';
+
+let isInsideBatchedSchedule = false;
+
+export function batch(fn) {
+  // if we are in node/tests or feature disabled or nested schedule
+  if (
+    !defaults.batchUpdates ||
+    !supports.scheduling() ||
+    isInsideBatchedSchedule
+  ) {
+    return unstable_batchedUpdates(fn);
+  }
+
+  isInsideBatchedSchedule = true;
+  // Use UserBlockingPriority as it has max 250ms timeout
+  // https://github.com/facebook/react/blob/master/packages/scheduler/src/forks/SchedulerNoDOM.js#L47
+  return scheduleCallback(
+    UserBlockingPriority,
+    function scheduleBatchedUpdates() {
+      unstable_batchedUpdates(fn);
+      isInsideBatchedSchedule = false;
+    }
+  );
+}

src/utils/batched-updates.native.js

@@ -1,2 +1,2 @@
 /* eslint-disable import/no-unresolved */
-export { unstable_batchedUpdates } from 'react-native';
+export { unstable_batchedUpdates as batch } from 'react-native';

src/utils/schedule.js

@@ -1,9 +1,5 @@
-import {
-  unstable_scheduleCallback as scheduleCallback,
-  unstable_UserBlockingPriority as UserBlockingPriority,
-} from 'scheduler';
 import defaults from '../defaults';
-import { unstable_batchedUpdates as batch } from './batched-updates';
+import { batch } from './batched-updates';
 import supports from './supported-features';
 
 const QUEUE = [];
@@ -20,14 +16,9 @@ export default function schedule(fn) {
 
   // if something already started schedule, skip
   if (scheduled) return;
-
-  // Use UserBlockingPriority as it has max 250ms timeout
-  // https://github.com/facebook/react/blob/master/packages/scheduler/src/forks/SchedulerNoDOM.js#L47
-  scheduled = scheduleCallback(UserBlockingPriority, function runNotifyQueue() {
-    batch(() => {
-      let queueFn;
-      while ((queueFn = QUEUE.shift())) queueFn();
-      scheduled = null;
-    });
+  scheduled = batch(() => {
+    let queueFn;
+    while ((queueFn = QUEUE.shift())) queueFn();
+    scheduled = null;
   });
 }