Rewrite subscription handling to batch actions

When a component mounts and tries to subscribe, it dispatches a thunk.
The first thunk dispatched for a given cache entry makes the actual
request. Later thunks for the same cache entry will see that the cache
entry exists, return `false` from `condition()`, and bail out.  That
dispatches a `rejected` action, and the reducer logic listened for that
to add the additional subscription entry.

However, this meant that adding N subscriptions to a cache entry would
dispatch N + 1: `pending` + `fulfilled` from the first thunk, and N-1
`rejected` actions from all others.

We've had reports that this led to slow perf when rendering large lists
of components with query hooks inside.

I've rewritten the logic to catch all `query/rejected` subscription
actions within an event loop tick, and dispatch them as a single
batched action at the end of the tick with `queueMicrotask`. This
drastically minimizes the number of dispatched actions - for a list of
1000 items, it went from 1002 actions to 4, and 6.5s to 1s.

Author: markerikson

packages/toolkit/src/query/core/buildMiddleware/batchActions.ts

@@ -0,0 +1,57 @@
+import type { QueryThunk, RejectedAction } from '../buildThunks'
+import type { SubMiddlewareBuilder } from './types'
+
+// 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)
+        )
+
+export const build: SubMiddlewareBuilder = ({
+  api,
+  context: { apiUid },
+  queryThunk,
+  reducerPath,
+}) => {
+  return (mwApi) => {
+    let abortedQueryActionsQueue: RejectedAction<QueryThunk, any>[] = []
+    let dispatchQueued = false
+
+    return (next) => (action) => {
+      if (queryThunk.rejected.match(action)) {
+        const { condition, arg } = action.meta
+
+        if (condition && arg.subscribe) {
+          // request was aborted due to condition (another query already running)
+          // _Don't_ dispatch right away - queue it for a debounced grouped dispatch
+          abortedQueryActionsQueue.push(action)
+
+          if (!dispatchQueued) {
+            queueMicrotaskShim(() => {
+              mwApi.dispatch(
+                api.internalActions.subscriptionRequestsRejected(
+                  abortedQueryActionsQueue
+                )
+              )
+              abortedQueryActionsQueue = []
+            })
+            dispatchQueued = true
+          }
+          // _Don't_ let the action reach the reducers now!
+          return
+        }
+      }
+
+      const result = next(action)
+
+      return result
+    }
+  }
+}

packages/toolkit/src/query/core/buildMiddleware/index.ts

@@ -17,6 +17,7 @@ import { build as buildWindowEventHandling } from './windowEventHandling'
 import { build as buildCacheLifecycle } from './cacheLifecycle'
 import { build as buildQueryLifecycle } from './queryLifecycle'
 import { build as buildDevMiddleware } from './devMiddleware'
+import { build as buildBatchActions } from './batchActions'
 
 export function buildMiddleware<
   Definitions extends EndpointDefinitions,
@@ -38,6 +39,7 @@ export function buildMiddleware<
     buildWindowEventHandling,
     buildCacheLifecycle,
     buildQueryLifecycle,
+    buildBatchActions,
   ].map((build) =>
     build({
       ...(input as any as BuildMiddlewareInput<

packages/toolkit/src/query/core/buildSlice.ts

@@ -23,7 +23,7 @@ import type {
   ConfigState,
 } from './apiState'
 import { QueryStatus } from './apiState'
-import type { MutationThunk, QueryThunk } from './buildThunks'
+import type { MutationThunk, QueryThunk, RejectedAction } from './buildThunks'
 import { calculateProvidedByThunk } from './buildThunks'
 import type {
   AssertTagTypes,
@@ -387,6 +387,26 @@ export function buildSlice({
           delete draft[queryCacheKey]![requestId]
         }
       },
+      subscriptionRequestsRejected(
+        draft,
+        action: PayloadAction<RejectedAction<QueryThunk, any>[]>
+      ) {
+        // We need to process "rejected" actions caused by a component trying to start a subscription
+        // after there's already a cache entry. Since many components may mount at once and all want
+        // the same data, we use a middleware that intercepts those actions batches these together
+        // into a single larger action , and we'll process all of them at once.
+        for (let rejectedAction of action.payload) {
+          const {
+            meta: { condition, arg, requestId },
+          } = rejectedAction
+          // request was aborted due to condition (another query already running)
+          if (condition && arg.subscribe) {
+            const substate = (draft[arg.queryCacheKey] ??= {})
+            substate[requestId] =
+              arg.subscriptionOptions ?? substate[requestId] ?? {}
+          }
+        }
+      },
     },
     extraReducers: (builder) => {
       builder
@@ -403,17 +423,6 @@ export function buildSlice({
               arg.subscriptionOptions ?? substate[requestId] ?? {}
           }
         })
-        .addCase(
-          queryThunk.rejected,
-          (draft, { meta: { condition, arg, requestId }, error, payload }) => {
-            // request was aborted due to condition (another query already running)
-            if (condition && arg.subscribe) {
-              const substate = (draft[arg.queryCacheKey] ??= {})
-              substate[requestId] =
-                arg.subscriptionOptions ?? substate[requestId] ?? {}
-            }
-          }
-        )
         // update the state to be a new object to be picked up as a "state change"
         // by redux-persist's `autoMergeLevel2`
         .addMatcher(hasRehydrationInfo, (draft) => ({ ...draft }))

packages/toolkit/src/query/react/buildHooks.ts

@@ -696,16 +696,25 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         pollingInterval,
       })
 
+      const lastRenderHadSubscription = useRef(false)
+
       const promiseRef = useRef<QueryActionCreatorResult<any>>()
 
       let { queryCacheKey, requestId } = promiseRef.current || {}
-      const subscriptionRemoved = useSelector(
+      const currentRenderHasSubscription = useSelector(
         (state: RootState<Definitions, string, string>) =>
           !!queryCacheKey &&
           !!requestId &&
           !state[api.reducerPath].subscriptions[queryCacheKey]?.[requestId]
       )
 
+      const subscriptionRemoved =
+        !currentRenderHasSubscription && lastRenderHadSubscription.current
+
+      usePossiblyImmediateEffect(() => {
+        lastRenderHadSubscription.current = currentRenderHasSubscription
+      })
+
       usePossiblyImmediateEffect((): void | undefined => {
         promiseRef.current = undefined
       }, [subscriptionRemoved])
@@ -736,6 +745,7 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
               forceRefetch: refetchOnMountOrArgChange,
             })
           )
+
           promiseRef.current = promise
         } else if (stableSubscriptionOptions !== lastSubscriptionOptions) {
           lastPromise.updateSubscriptionOptions(stableSubscriptionOptions)
@@ -923,8 +933,9 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
           ...options,
         })
 
-        const { data, status, isLoading, isSuccess, isError, error } = queryStateResults;
-        useDebugValue({ data, status, isLoading, isSuccess, isError, error });
+        const { data, status, isLoading, isSuccess, isError, error } =
+          queryStateResults
+        useDebugValue({ data, status, isLoading, isSuccess, isError, error })
 
         return useMemo(
           () => ({ ...queryStateResults, ...querySubscriptionResults }),
@@ -993,8 +1004,24 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         })
       }, [dispatch, fixedCacheKey, promise, requestId])
 
-      const { endpointName, data, status, isLoading, isSuccess, isError, error } = currentState;
-      useDebugValue({ endpointName, data, status, isLoading, isSuccess, isError, error });
+      const {
+        endpointName,
+        data,
+        status,
+        isLoading,
+        isSuccess,
+        isError,
+        error,
+      } = currentState
+      useDebugValue({
+        endpointName,
+        data,
+        status,
+        isLoading,
+        isSuccess,
+        isError,
+        error,
+      })
 
       const finalState = useMemo(
         () => ({ ...currentState, originalArgs, reset }),

packages/toolkit/src/query/tests/buildHooks.test.tsx

@@ -1162,9 +1162,12 @@ describe('hooks tests', () => {
     })
 
     test('useMutation return value contains originalArgs', async () => {
-      const { result } = renderHook(() => api.endpoints.updateUser.useMutation(), {
-        wrapper: storeRef.wrapper,
-      })
+      const { result } = renderHook(
+        () => api.endpoints.updateUser.useMutation(),
+        {
+          wrapper: storeRef.wrapper,
+        }
+      )
       const arg = { name: 'Foo' }
 
       const firstRenderResult = result.current
@@ -1955,13 +1958,13 @@ describe('hooks with createApi defaults set', () => {
 
       const addBtn = screen.getByTestId('addPost')
 
-      await waitFor(() => expect(getRenderCount()).toBe(3))
+      await waitFor(() => expect(getRenderCount()).toBe(4))
 
       fireEvent.click(addBtn)
-      await waitFor(() => expect(getRenderCount()).toBe(5))
+      await waitFor(() => expect(getRenderCount()).toBe(6))
       fireEvent.click(addBtn)
       fireEvent.click(addBtn)
-      await waitFor(() => expect(getRenderCount()).toBe(7))
+      await waitFor(() => expect(getRenderCount()).toBe(8))
     })
 
     test('useQuery with selectFromResult option serves a deeply memoized value and does not rerender unnecessarily', async () => {

packages/toolkit/src/query/tests/cleanup.test.tsx

@@ -2,12 +2,13 @@
 
 import React from 'react'
 
+import { createListenerMiddleware } from '@reduxjs/toolkit'
 import { createApi, QueryStatus } from '@reduxjs/toolkit/query/react'
-import { render, waitFor } from '@testing-library/react'
+import { render, waitFor, act, screen } from '@testing-library/react'
 import { setupApiStore } from './helpers'
 
 const api = createApi({
-  baseQuery: () => ({ data: null }),
+  baseQuery: () => ({ data: 42 }),
   endpoints: (build) => ({
     a: build.query<unknown, void>({ query: () => '' }),
     b: build.query<unknown, void>({ query: () => '' }),
@@ -19,18 +20,21 @@ let getSubStateA = () => storeRef.store.getState().api.queries['a(undefined)']
 let getSubStateB = () => storeRef.store.getState().api.queries['b(undefined)']
 
 function UsingA() {
-  api.endpoints.a.useQuery()
-  return <></>
+  const { data } = api.endpoints.a.useQuery()
+
+  return <>Result: {data} </>
 }
 
 function UsingB() {
   api.endpoints.b.useQuery()
+
   return <></>
 }
 
 function UsingAB() {
   api.endpoints.a.useQuery()
   api.endpoints.b.useQuery()
+
   return <></>
 }
 
@@ -90,11 +94,15 @@ test('data stays in store when component stays rendered while data for another c
 
   const statusA = getSubStateA()
 
-  rerender(
-    <>
-      <UsingA />
-    </>
-  )
+  await act(async () => {
+    rerender(
+      <>
+        <UsingA />
+      </>
+    )
+
+    jest.advanceTimersByTime(10)
+  })
 
   jest.advanceTimersByTime(120000)
 
@@ -108,8 +116,8 @@ test('data stays in store when one component requiring the data stays in the sto
 
   const { rerender } = render(
     <>
-      <UsingA />
-      <UsingAB />
+      <UsingA key="a" />
+      <UsingAB key="ab" />
     </>,
     { wrapper: storeRef.wrapper }
   )
@@ -121,14 +129,81 @@ test('data stays in store when one component requiring the data stays in the sto
   const statusA = getSubStateA()
   const statusB = getSubStateB()
 
-  rerender(
-    <>
-      <UsingAB />
-    </>
-  )
+  await act(async () => {
+    rerender(
+      <>
+        <UsingAB key="ab" />
+      </>
+    )
+    jest.advanceTimersByTime(10)
+    jest.runAllTimers()
+  })
 
-  jest.advanceTimersByTime(120000)
+  await act(async () => {
+    jest.advanceTimersByTime(120000)
+    jest.runAllTimers()
+  })
 
   expect(getSubStateA()).toEqual(statusA)
   expect(getSubStateB()).toEqual(statusB)
 })
+
+test('Minimizes the number of subscription dispatches when multiple components ask for the same data', async () => {
+  const listenerMiddleware = createListenerMiddleware()
+  const storeRef = setupApiStore(api, undefined, {
+    middleware: [listenerMiddleware.middleware],
+    withoutTestLifecycles: true,
+  })
+
+  let getSubscriptionsA = () =>
+    storeRef.store.getState().api.subscriptions['a(undefined)']
+
+  let actionTypes: string[] = []
+
+  listenerMiddleware.startListening({
+    predicate: () => true,
+    effect: (action) => {
+      actionTypes.push(action.type)
+    },
+  })
+
+  const NUM_LIST_ITEMS = 1000
+
+  function ParentComponent() {
+    const listItems = Array.from({ length: NUM_LIST_ITEMS }).map((_, i) => (
+      <UsingA key={i} />
+    ))
+
+    return <>{listItems}</>
+  }
+
+  const start = Date.now()
+
+  render(<ParentComponent />, {
+    wrapper: storeRef.wrapper,
+  })
+
+  jest.advanceTimersByTime(10)
+
+  await waitFor(() => {
+    return screen.getAllByText(/42/).length > 0
+  })
+
+  const end = Date.now()
+
+  const timeElapsed = end - start
+
+  const subscriptions = getSubscriptionsA()
+
+  expect(Object.keys(subscriptions!).length).toBe(NUM_LIST_ITEMS)
+  // Expected: [
+  //   'api/config/middlewareRegistered',
+  //   'api/executeQuery/pending',
+  //   'api/subscriptions/subscriptionRequestsRejected',
+  //   'api/executeQuery/fulfilled'
+  // ]
+  expect(actionTypes.length).toBe(4)
+  // Could be flaky in CI, but we'll see.
+  // Currently seeing 1000ms in local dev, 6300 without the batching fixes
+  expect(timeElapsed).toBeLessThan(2500)
+}, 25000)