Add an auto-batching enhancer that delays low-pri notifications

Author: markerikson

packages/toolkit/src/autoBatchEnhancer.ts

@@ -0,0 +1,110 @@
+import type { StoreEnhancer } from 'redux'
+
+export const SHOULD_AUTOBATCH = 'RTK_autoBatch'
+
+export const prepareAutoBatched =
+  <T>() =>
+  (payload: T): { payload: T; meta: unknown } => ({
+    payload,
+    meta: { [SHOULD_AUTOBATCH]: true },
+  })
+
+// TODO Remove this in 2.0
+// Copied from https://github.com/feross/queue-microtask
+let promise: Promise<any>
+const queueMicrotaskShim =
+  typeof queueMicrotask === 'function'
+    ? queueMicrotask.bind(typeof window !== 'undefined' ? window : global)
+    : // reuse resolved promise, and allocate it lazily
+      (cb: () => void) =>
+        (promise || (promise = Promise.resolve())).then(cb).catch((err: any) =>
+          setTimeout(() => {
+            throw err
+          }, 0)
+        )
+
+/**
+ * A Redux store enhancer that watches for "low-priority" actions, and delays
+ * notifying subscribers until either the end of the event loop tick or the
+ * next "standard-priority" action is dispatched.
+ *
+ * This allows dispatching multiple "low-priority" actions in a row with only
+ * a single subscriber notification to the UI after the sequence of actions
+ * is finished, thus improving UI re-render performance.
+ *
+ * Watches for actions with the `action.meta[SHOULD_AUTOBATCH]` attribute.
+ * This can be added to `action.meta` manually, or by using the
+ * `prepareAutoBatched` helper.
+ *
+ */
+export const autoBatchEnhancer =
+  (): StoreEnhancer =>
+  (next) =>
+  (...args) => {
+    const store = next(...args)
+
+    let notifying = true
+    let shouldNotifyAtEndOfTick = false
+    let notificationQueued = false
+
+    const listeners = new Set<() => void>()
+
+    const notifyListeners = () => {
+      // We're running at the end of the event loop tick.
+      // Run the real listener callbacks to actually update the UI.
+      notificationQueued = false
+      if (shouldNotifyAtEndOfTick) {
+        shouldNotifyAtEndOfTick = false
+        listeners.forEach((l) => l())
+      }
+    }
+
+    return Object.assign({}, store, {
+      // Override the base `store.subscribe` method to keep original listeners
+      // from running if we're delaying notifications
+      subscribe(listener: () => void) {
+        // Each wrapped listener will only call the real listener if
+        // the `notifying` flag is currently active when it's called.
+        // This lets the base store work as normal, while the actual UI
+        // update becomes controlled by this enhancer.
+        const wrappedListener: typeof listener = () => notifying && listener()
+        const unsubscribe = store.subscribe(wrappedListener)
+        listeners.add(listener)
+        return () => {
+          unsubscribe()
+          listeners.delete(listener)
+        }
+      },
+      // Override the base `store.dispatch` method so that we can check actions
+      // for the `shouldAutoBatch` flag and determine if batching is active
+      dispatch(action: any) {
+        try {
+          // If the action does _not_ have the `shouldAutoBatch` flag,
+          // we resume/continue normal notify-after-each-dispatch behavior
+          notifying = !action?.meta?.[SHOULD_AUTOBATCH]
+          // If a `notifyListeners` microtask was queued, you can't cancel it.
+          // Instead, we set a flag so that it's a no-op when it does run
+          shouldNotifyAtEndOfTick = !notifying
+          if (shouldNotifyAtEndOfTick) {
+            // We've seen at least 1 action with `SHOULD_AUTOBATCH`. Try to queue
+            // a microtask to notify listeners at the end of the event loop tick.
+            // Make sure we only enqueue this _once_ per tick.
+            if (!notificationQueued) {
+              notificationQueued = true
+              queueMicrotaskShim(notifyListeners)
+            }
+          }
+          // Go ahead and process the action as usual, including reducers.
+          // If normal notification behavior is enabled, the store will notify
+          // all of its own listeners, and the wrapper callbacks above will
+          // see `notifying` is true and pass on to the real listener callbacks.
+          // If we're "batching" behavior, then the wrapped callbacks will
+          // bail out, causing the base store notification behavior to be no-ops.
+          return store.dispatch(action)
+        } finally {
+          // Assume we're back to normal behavior after each action
+          notifying = true
+        }
+      },
+    })
+  }

packages/toolkit/src/index.ts

@@ -184,3 +184,9 @@ export {
   clearAllListeners,
   TaskAbortError,
 } from './listenerMiddleware/index'
+
+export {
+  SHOULD_AUTOBATCH,
+  prepareAutoBatched,
+  autoBatchEnhancer,
+} from './autoBatchEnhancer'

packages/toolkit/src/tests/autoBatchEnhancer.test.ts

@@ -0,0 +1,111 @@
+import { configureStore } from '../configureStore'
+import { createSlice } from '../createSlice'
+import { autoBatchEnhancer, prepareAutoBatched } from '../autoBatchEnhancer'
+import { delay } from '../utils'
+
+interface CounterState {
+  value: number
+}
+
+const counterSlice = createSlice({
+  name: 'counter',
+  initialState: { value: 0 } as CounterState,
+  reducers: {
+    incrementBatched: {
+      // Batched, low-priority
+      reducer(state) {
+        state.value += 1
+      },
+      prepare: prepareAutoBatched<void>(),
+    },
+    // Not batched, normal priority
+    decrementUnbatched(state) {
+      state.value -= 1
+    },
+  },
+})
+const { incrementBatched, decrementUnbatched } = counterSlice.actions
+
+const makeStore = () => {
+  return configureStore({
+    reducer: counterSlice.reducer,
+    enhancers: (existingEnhancers) => {
+      return existingEnhancers.concat(autoBatchEnhancer())
+    },
+  })
+}
+
+let store: ReturnType<typeof makeStore>
+
+let subscriptionNotifications = 0
+
+beforeEach(() => {
+  subscriptionNotifications = 0
+  store = makeStore()
+
+  store.subscribe(() => {
+    subscriptionNotifications++
+  })
+})
+
+describe('autoBatchEnhancer', () => {
+  test('Does not alter normal subscription notification behavior', async () => {
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(1)
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(2)
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(3)
+    store.dispatch(decrementUnbatched())
+
+    await delay(5)
+
+    expect(subscriptionNotifications).toBe(4)
+  })
+
+  test('Only notifies once if several batched actions are dispatched in a row', async () => {
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(incrementBatched())
+
+    await delay(5)
+
+    expect(subscriptionNotifications).toBe(1)
+  })
+
+  test('Notifies immediately if a non-batched action is dispatched', async () => {
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(1)
+    store.dispatch(incrementBatched())
+
+    await delay(5)
+
+    expect(subscriptionNotifications).toBe(2)
+  })
+
+  test('Does not notify at end of tick if last action was normal priority', async () => {
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(incrementBatched())
+    expect(subscriptionNotifications).toBe(0)
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(1)
+    store.dispatch(incrementBatched())
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(2)
+    store.dispatch(decrementUnbatched())
+    expect(subscriptionNotifications).toBe(3)
+
+    await delay(5)
+
+    expect(subscriptionNotifications).toBe(3)
+  })
+})